Skip to content
Closed
Show file tree
Hide file tree
Changes from all commits
Commits
File filter

Filter by extension

Filter by extension

Conversations
Failed to load comments.
Loading
Jump to
Jump to file
Failed to load files.
Loading
Diff view
Diff view
221 changes: 221 additions & 0 deletions cmd/restore.go
Original file line number Diff line number Diff line change
@@ -0,0 +1,221 @@
package cmd

import (
"archive/zip"
"context"
"fmt"
"io"
"os"
"os/exec"
"path/filepath"
"time"

"github.com/spf13/cobra"

"github.com/supermodeltools/cli/internal/api"
"github.com/supermodeltools/cli/internal/config"
"github.com/supermodeltools/cli/internal/restore"
)

func init() {
var localMode bool
var maxTokens int
var dir string

c := &cobra.Command{
Use: "restore",
Short: "Generate a project context summary to restore Claude's understanding",
Long: `Restore builds a high-level project summary (a "context bomb") and writes it
to stdout. Use it after Claude Code compacts its context window to re-establish
understanding of your codebase structure, domains, and key files.

With an API key configured (run 'supermodel login'), restore calls the
Supermodel API for an AI-powered analysis including semantic domains, external
dependencies, and critical file ranking.

Without an API key (or with --local), restore performs a local scan of the
repository file tree and produces a simpler structural summary.

Examples:

# pipe into Claude Code (typical use)
supermodel restore

# use local analysis only, no API call
supermodel restore --local

# increase the token budget for larger projects
supermodel restore --max-tokens 4000`,
RunE: func(cmd *cobra.Command, args []string) error {
return runRestore(cmd, dir, localMode, maxTokens)
},
SilenceUsage: true,
}

c.Flags().BoolVar(&localMode, "local", false, "use local file scan instead of Supermodel API")
c.Flags().IntVar(&maxTokens, "max-tokens", restore.DefaultMaxTokens, "maximum token budget for the output")
c.Flags().StringVar(&dir, "dir", "", "project directory (default: current working directory)")

rootCmd.AddCommand(c)
}

func runRestore(cmd *cobra.Command, dir string, localMode bool, maxTokens int) error {
// Resolve the project directory.
if dir == "" {
var err error
dir, err = os.Getwd()
if err != nil {
return fmt.Errorf("get working directory: %w", err)
}
}
rootDir := findGitRoot(dir)

projectName := filepath.Base(rootDir)

opts := restore.RenderOptions{
MaxTokens: maxTokens,
ClaudeMD: restore.ReadClaudeMD(rootDir),
}

var graph *restore.ProjectGraph

cfg, _ := config.Load()
hasAPIKey := cfg != nil && cfg.APIKey != ""

if !localMode && hasAPIKey {
var err error
graph, err = restoreViaAPI(cmd, cfg, rootDir, projectName)
if err != nil {
fmt.Fprintf(cmd.ErrOrStderr(), "warning: API analysis failed (%v), falling back to local mode\n", err)
}
}

if graph == nil {
opts.LocalMode = true
ctx, cancel := context.WithTimeout(context.Background(), 30*time.Second)
defer cancel()
var err error
graph, err = restore.BuildProjectGraph(ctx, rootDir, projectName)
if err != nil {
return fmt.Errorf("local analysis failed: %w", err)
}
}

output, _, err := restore.Render(graph, projectName, opts)
if err != nil {
return fmt.Errorf("render: %w", err)
}
_, err = fmt.Fprint(cmd.OutOrStdout(), output)
return err
}

func restoreViaAPI(cmd *cobra.Command, cfg *config.Config, rootDir, projectName string) (*restore.ProjectGraph, error) {
zipPath, err := restoreCreateZip(rootDir)
if err != nil {
return nil, fmt.Errorf("create archive: %w", err)
}
defer os.Remove(zipPath)

client := api.New(cfg)

ctx, cancel := context.WithTimeout(context.Background(), 10*time.Minute)
defer cancel()

fmt.Fprintln(cmd.ErrOrStderr(), "Analyzing repository…")
ir, err := client.AnalyzeDomains(ctx, zipPath, "restore-"+projectName)
if err != nil {
return nil, err
}

graph := restore.FromSupermodelIR(ir, projectName)
return graph, nil
}

// restoreCreateZip creates a temporary ZIP of the repository at dir.
// It tries git archive first (respects .gitignore), then falls back to a
// simple directory walk.
func restoreCreateZip(dir string) (string, error) {
f, err := os.CreateTemp("", "supermodel-restore-*.zip")
if err != nil {
return "", err
}
dest := f.Name()
f.Close()

cmd := exec.Command("git", "-C", dir, "archive", "--format=zip", "-o", dest, "HEAD")
cmd.Stderr = os.Stderr
if err := cmd.Run(); err == nil {
return dest, nil
}

// Fallback: walk the directory.
if err := restoreWalkZip(dir, dest); err != nil {
_ = os.Remove(dest)
return "", err
}
return dest, nil
}

// restoreWalkZip archives dir into a ZIP at dest, skipping common build/cache dirs.
func restoreWalkZip(dir, dest string) error {
out, err := os.Create(dest) //nolint:gosec // dest is a temp file path from os.CreateTemp
if err != nil {
return err
}
defer out.Close()

zw := zip.NewWriter(out)
defer zw.Close()

skipDirs := map[string]bool{
".git": true, "node_modules": true, "vendor": true, "__pycache__": true,
".venv": true, "venv": true, "dist": true, "build": true, "target": true,
".next": true, ".nuxt": true, "coverage": true, ".terraform": true,
}

return filepath.Walk(dir, func(path string, info os.FileInfo, err error) error {
if err != nil {
return nil
}
rel, err := filepath.Rel(dir, path)
if err != nil {
return nil
}
if info.IsDir() {
if skipDirs[info.Name()] {
return filepath.SkipDir
}
return nil
}
if info.Size() > 10<<20 {
return nil
}
w, err := zw.Create(filepath.ToSlash(rel))
if err != nil {
return err
}
src, err := os.Open(path) //nolint:gosec // path is from filepath.Walk within dir
if err != nil {
return err
}
defer src.Close()
_, err = io.Copy(w, src)
return err
})
}

// findGitRoot walks up from start to find the directory containing .git.
// Returns start itself if no .git directory is found.
func findGitRoot(start string) string {
dir := start
for {
if _, err := os.Stat(filepath.Join(dir, ".git")); err == nil {
return dir
}
parent := filepath.Dir(dir)
if parent == dir {
return start
}
dir = parent
}
}
40 changes: 30 additions & 10 deletions internal/api/client.go
Original file line number Diff line number Diff line change
Expand Up @@ -39,12 +39,39 @@ const analyzeEndpoint = "/v1/graphs/supermodel"
// Analyze uploads a repository ZIP and runs the full analysis pipeline,
// polling until the async job completes and returning the Graph.
func (c *Client) Analyze(ctx context.Context, zipPath, idempotencyKey string) (*Graph, error) {
job, err := c.postZip(ctx, zipPath, idempotencyKey)
job, err := c.pollUntilComplete(ctx, zipPath, idempotencyKey)
if err != nil {
return nil, err
}
var result jobResult
if err := json.Unmarshal(job.Result, &result); err != nil {
return nil, fmt.Errorf("decode graph result: %w", err)
}
return &result.Graph, nil
}

// AnalyzeDomains uploads a repository ZIP and runs the full analysis pipeline,
// returning the complete SupermodelIR response (domains, summary, metadata, graph).
// Use this instead of Analyze when you need high-level domain information.
func (c *Client) AnalyzeDomains(ctx context.Context, zipPath, idempotencyKey string) (*SupermodelIR, error) {
job, err := c.pollUntilComplete(ctx, zipPath, idempotencyKey)
if err != nil {
return nil, err
}
var ir SupermodelIR
if err := json.Unmarshal(job.Result, &ir); err != nil {
return nil, fmt.Errorf("decode domain result: %w", err)
}
return &ir, nil
}

// Poll until the job completes.
// pollUntilComplete submits a ZIP to the analyze endpoint and polls until the
// async job reaches "completed" status, then returns the raw JobResponse.
func (c *Client) pollUntilComplete(ctx context.Context, zipPath, idempotencyKey string) (*JobResponse, error) {
job, err := c.postZip(ctx, zipPath, idempotencyKey)
if err != nil {
return nil, err
}
for job.Status == "pending" || job.Status == "processing" {
wait := time.Duration(job.RetryAfter) * time.Second
if wait <= 0 {
Expand All @@ -55,25 +82,18 @@ func (c *Client) Analyze(ctx context.Context, zipPath, idempotencyKey string) (*
return nil, ctx.Err()
case <-time.After(wait):
}

job, err = c.postZip(ctx, zipPath, idempotencyKey)
if err != nil {
return nil, err
}
}

if job.Error != nil {
return nil, fmt.Errorf("analysis failed: %s", *job.Error)
}
if job.Status != "completed" {
return nil, fmt.Errorf("unexpected job status: %s", job.Status)
}

var result jobResult
if err := json.Unmarshal(job.Result, &result); err != nil {
return nil, fmt.Errorf("decode graph result: %w", err)
}
return &result.Graph, nil
return job, nil
}

// postZip sends the repository ZIP to the analyze endpoint and returns the
Expand Down
51 changes: 51 additions & 0 deletions internal/api/types.go
Original file line number Diff line number Diff line change
Expand Up @@ -93,6 +93,57 @@ func (g *Graph) NodeByID(id string) (Node, bool) {
return Node{}, false
}

// SupermodelIR is the full structured response returned inside a completed job
// result from /v1/graphs/supermodel. It contains high-level domain information
// in addition to the raw node/edge graph captured by Graph.
type SupermodelIR struct {
Repo string `json:"repo"`
Summary map[string]any `json:"summary"`
Metadata IRMetadata `json:"metadata"`
Domains []IRDomain `json:"domains"`
Graph IRGraph `json:"graph"`
}

// IRMetadata holds file-count and language statistics from the API response.
type IRMetadata struct {
FileCount int `json:"fileCount"`
Languages []string `json:"languages"`
}

// IRGraph is the raw node/relationship sub-graph embedded in SupermodelIR.
type IRGraph struct {
Nodes []IRNode `json:"nodes"`
Relationships []IRRelationship `json:"relationships"`
}

// IRNode is a single node in the IRGraph.
type IRNode struct {
Type string `json:"type"`
Name string `json:"name"`
}

// IRRelationship is a directed edge in the IRGraph.
type IRRelationship struct {
Type string `json:"type"`
Source string `json:"source"`
Target string `json:"target"`
}

// IRDomain is the raw representation of a semantic domain from the API.
type IRDomain struct {
Name string `json:"name"`
DescriptionSummary string `json:"descriptionSummary"`
KeyFiles []string `json:"keyFiles"`
Responsibilities []string `json:"responsibilities"`
Subdomains []IRSubdomain `json:"subdomains"`
}

// IRSubdomain is a named sub-area within an IRDomain.
type IRSubdomain struct {
Name string `json:"name"`
DescriptionSummary string `json:"descriptionSummary"`
}

// JobResponse is the async envelope returned by the API for long-running jobs.
type JobResponse struct {
Status string `json:"status"`
Expand Down
2 changes: 1 addition & 1 deletion internal/blastradius/handler.go
Original file line number Diff line number Diff line change
Expand Up @@ -62,7 +62,7 @@ func findBlastRadius(g *api.Graph, repoDir, target string, maxDepth int) ([]Resu
// Build reverse adjacency: nodeID → set of node IDs that import it.
importedBy := make(map[string][]string)
for _, rel := range g.Rels() {
if rel.Type == "IMPORTS" || rel.Type == "WILDCARD_IMPORTS" {
if rel.Type == "imports" || rel.Type == "wildcard_imports" {
importedBy[rel.EndNode] = append(importedBy[rel.EndNode], rel.StartNode)
}
}
Expand Down
4 changes: 2 additions & 2 deletions internal/blastradius/handler_test.go
Original file line number Diff line number Diff line change
Expand Up @@ -37,8 +37,8 @@ func TestFindBlastRadius(t *testing.T) {
fileNode("c", "internal/c/c.go"),
},
Relationships: []api.Relationship{
{ID: "r1", Type: "IMPORTS", StartNode: "a", EndNode: "b"},
{ID: "r2", Type: "IMPORTS", StartNode: "c", EndNode: "a"},
{ID: "r1", Type: "imports", StartNode: "a", EndNode: "b"},
{ID: "r2", Type: "imports", StartNode: "c", EndNode: "a"},
},
}

Expand Down
2 changes: 1 addition & 1 deletion internal/deadcode/handler.go
Original file line number Diff line number Diff line change
Expand Up @@ -45,7 +45,7 @@ func findDeadCode(g *api.Graph, includeExports bool) []Result {
// Build set of function node IDs that receive at least one CALLS edge.
called := make(map[string]bool)
for _, rel := range g.Rels() {
if rel.Type == "CALLS" || rel.Type == "CONTAINS_CALL" {
if rel.Type == "calls" || rel.Type == "contains_call" {
called[rel.EndNode] = true
}
}
Expand Down
2 changes: 1 addition & 1 deletion internal/deadcode/handler_test.go
Original file line number Diff line number Diff line change
Expand Up @@ -46,7 +46,7 @@ func TestFindDeadCode(t *testing.T) {
node("f4", "Function", "Handler", "server.go"), // exported → excluded by default
},
Relationships: []api.Relationship{
{ID: "r1", Type: "CALLS", StartNode: "f2", EndNode: "f1"}, // main → process
{ID: "r1", Type: "calls", StartNode: "f2", EndNode: "f1"}, // main → process
},
}

Expand Down
Loading
Loading