Skip to content

Auto-correct Go imports in Code Mode #118

New issue

Have a question about this project? Sign up for a free GitHub account to open an issue and contact its maintainers and the community.

Sign up for GitHub

By clicking “Sign up for GitHub”, you agree to our terms of service and privacy statement. We’ll occasionally send you account related emails.

Already on GitHub? Sign in to your account

Merged
spachava753 merged 3 commits into from
Dec 20, 2025
Merged

Auto-correct Go imports in Code Mode #118

Show file tree
Hide file tree
Changes from all commits
Commits
Show all changes
3 commits
Select commit Hold shift + click to select a range
1dea8b9
feat(codemode): auto-correct Go imports using goimports
spachava753 Dec 20, 2025
4521281
refactor(codemode): use x/tools/imports instead of goimports CLI
spachava753 Dec 20, 2025
c10eb46
feat(codemode): list added/removed packages in auto-correct message
spachava753 Dec 20, 2025
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
14 changes: 8 additions & 6 deletions go.mod
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -29,7 +29,7 @@ require (
github.com/spf13/cobra v1.10.1
github.com/stoewer/go-strcase v1.2.0
github.com/stretchr/testify v1.11.1
golang.org/x/term v0.37.0
golang.org/x/term v0.38.0
google.golang.org/genai v1.36.0
gopkg.in/yaml.v3 v3.0.1
)
Expand Down Expand Up @@ -125,13 +125,15 @@ require (
go.uber.org/atomic v1.11.0 // indirect
go.uber.org/multierr v1.11.0 // indirect
go.uber.org/zap v1.27.0 // indirect
golang.org/x/crypto v0.44.0 // indirect
golang.org/x/crypto v0.46.0 // indirect
golang.org/x/exp v0.0.0-20240506185415-9bf2ced13842 // indirect
golang.org/x/net v0.47.0 // indirect
golang.org/x/mod v0.31.0 // indirect
golang.org/x/net v0.48.0 // indirect
golang.org/x/oauth2 v0.33.0 // indirect
golang.org/x/sync v0.18.0 // indirect
golang.org/x/sys v0.38.0 // indirect
golang.org/x/text v0.31.0 // indirect
golang.org/x/sync v0.19.0 // indirect
golang.org/x/sys v0.39.0 // indirect
golang.org/x/text v0.32.0 // indirect
golang.org/x/tools v0.40.0 // indirect
google.golang.org/genproto/googleapis/api v0.0.0-20251022142026-3a174f9686a8 // indirect
google.golang.org/genproto/googleapis/rpc v0.0.0-20251111163417-95abcf5c77ba // indirect
google.golang.org/grpc v1.77.0 // indirect
Expand Down
16 changes: 16 additions & 0 deletions go.sum
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -287,34 +287,50 @@ go.uber.org/zap v1.27.0/go.mod h1:GB2qFLM7cTU87MWRP2mPIjqfIDnGu+VIO4V/SdhGo2E=
golang.org/x/crypto v0.0.0-20190308221718-c2843e01d9a2/go.mod h1:djNgcEr1/C05ACkg1iLfiJU5Ep61QUkGW8qpdssI0+w=
golang.org/x/crypto v0.44.0 h1:A97SsFvM3AIwEEmTBiaxPPTYpDC47w720rdiiUvgoAU=
golang.org/x/crypto v0.44.0/go.mod h1:013i+Nw79BMiQiMsOPcVCB5ZIJbYkerPrGnOa00tvmc=
golang.org/x/crypto v0.46.0 h1:cKRW/pmt1pKAfetfu+RCEvjvZkA9RimPbh7bhFjGVBU=
golang.org/x/crypto v0.46.0/go.mod h1:Evb/oLKmMraqjZ2iQTwDwvCtJkczlDuTmdJXoZVzqU0=
golang.org/x/exp v0.0.0-20240506185415-9bf2ced13842 h1:vr/HnozRka3pE4EsMEg1lgkXJkTFJCVUX+S/ZT6wYzM=
golang.org/x/exp v0.0.0-20240506185415-9bf2ced13842/go.mod h1:XtvwrStGgqGPLc4cjQfWqZHG1YFdYs6swckp8vpsjnc=
golang.org/x/lint v0.0.0-20190930215403-16217165b5de/go.mod h1:6SW0HCj/g11FgYtHlgUYUwCkIfeOF89ocIRzGO/8vkc=
golang.org/x/mod v0.29.0 h1:HV8lRxZC4l2cr3Zq1LvtOsi/ThTgWnUk/y64QSs8GwA=
golang.org/x/mod v0.29.0/go.mod h1:NyhrlYXJ2H4eJiRy/WDBO6HMqZQ6q9nk4JzS3NuCK+w=
golang.org/x/mod v0.31.0 h1:HaW9xtz0+kOcWKwli0ZXy79Ix+UW/vOfmWI5QVd2tgI=
golang.org/x/mod v0.31.0/go.mod h1:43JraMp9cGx1Rx3AqioxrbrhNsLl2l/iNAvuBkrezpg=
golang.org/x/net v0.0.0-20190311183353-d8887717615a/go.mod h1:t9HGtf8HONx5eT2rtn7q6eTqICYqUVnKs3thJo3Qplg=
golang.org/x/net v0.0.0-20190620200207-3b0461eec859/go.mod h1:z5CRVTTTmAJ677TzLLGU+0bjPO0LkuOLi4/5GtJWs/s=
golang.org/x/net v0.47.0 h1:Mx+4dIFzqraBXUugkia1OOvlD6LemFo1ALMHjrXDOhY=
golang.org/x/net v0.47.0/go.mod h1:/jNxtkgq5yWUGYkaZGqo27cfGZ1c5Nen03aYrrKpVRU=
golang.org/x/net v0.48.0 h1:zyQRTTrjc33Lhh0fBgT/H3oZq9WuvRR5gPC70xpDiQU=
golang.org/x/net v0.48.0/go.mod h1:+ndRgGjkh8FGtu1w1FGbEC31if4VrNVMuKTgcAAnQRY=
golang.org/x/oauth2 v0.33.0 h1:4Q+qn+E5z8gPRJfmRy7C2gGG3T4jIprK6aSYgTXGRpo=
golang.org/x/oauth2 v0.33.0/go.mod h1:lzm5WQJQwKZ3nwavOZ3IS5Aulzxi68dUSgRHujetwEA=
golang.org/x/sync v0.0.0-20190423024810-112230192c58/go.mod h1:RxMgew5VJxzue5/jJTE5uejpjVlOe/izrB70Jof72aM=
golang.org/x/sync v0.18.0 h1:kr88TuHDroi+UVf+0hZnirlk8o8T+4MrK6mr60WkH/I=
golang.org/x/sync v0.18.0/go.mod h1:9KTHXmSnoGruLpwFjVSX0lNNA75CykiMECbovNTZqGI=
golang.org/x/sync v0.19.0 h1:vV+1eWNmZ5geRlYjzm2adRgW2/mcpevXNg50YZtPCE4=
golang.org/x/sync v0.19.0/go.mod h1:9KTHXmSnoGruLpwFjVSX0lNNA75CykiMECbovNTZqGI=
golang.org/x/sys v0.0.0-20190215142949-d0b11bdaac8a/go.mod h1:STP8DvDyc/dI5b8T5hshtkjS+E42TnysNCUPdjciGhY=
golang.org/x/sys v0.6.0/go.mod h1:oPkhp1MJrh7nUepCBck5+mAzfO9JrbApNNgaTdGDITg=
golang.org/x/sys v0.38.0 h1:3yZWxaJjBmCWXqhN1qh02AkOnCQ1poK6oF+a7xWL6Gc=
golang.org/x/sys v0.38.0/go.mod h1:OgkHotnGiDImocRcuBABYBEXf8A9a87e/uXjp9XT3ks=
golang.org/x/sys v0.39.0 h1:CvCKL8MeisomCi6qNZ+wbb0DN9E5AATixKsvNtMoMFk=
golang.org/x/sys v0.39.0/go.mod h1:OgkHotnGiDImocRcuBABYBEXf8A9a87e/uXjp9XT3ks=
golang.org/x/term v0.37.0 h1:8EGAD0qCmHYZg6J17DvsMy9/wJ7/D/4pV/wfnld5lTU=
golang.org/x/term v0.37.0/go.mod h1:5pB4lxRNYYVZuTLmy8oR2BH8dflOR+IbTYFD8fi3254=
golang.org/x/term v0.38.0 h1:PQ5pkm/rLO6HnxFR7N2lJHOZX6Kez5Y1gDSJla6jo7Q=
golang.org/x/term v0.38.0/go.mod h1:bSEAKrOT1W+VSu9TSCMtoGEOUcKxOKgl3LE5QEF/xVg=
golang.org/x/text v0.3.0/go.mod h1:NqM8EUOU14njkJ3fqMW+pc6Ldnwhi/IjpwHt7yyuwOQ=
golang.org/x/text v0.31.0 h1:aC8ghyu4JhP8VojJ2lEHBnochRno1sgL6nEi9WGFGMM=
golang.org/x/text v0.31.0/go.mod h1:tKRAlv61yKIjGGHX/4tP1LTbc13YSec1pxVEWXzfoeM=
golang.org/x/text v0.32.0 h1:ZD01bjUt1FQ9WJ0ClOL5vxgxOI/sVCNgX1YtKwcY0mU=
golang.org/x/text v0.32.0/go.mod h1:o/rUWzghvpD5TXrTIBuJU77MTaN0ljMWE47kxGJQ7jY=
golang.org/x/tools v0.0.0-20190311212946-11955173bddd/go.mod h1:LCzVGOaR6xXOjkQ3onu1FJEFr0SW1gC7cKk1uF8kGRs=
golang.org/x/tools v0.0.0-20191029041327-9cc4af7d6b2c/go.mod h1:b+2E5dAYhXwXZwtnZ6UAqBI28+e2cm9otk0dWdXHAEo=
golang.org/x/tools v0.0.0-20191108193012-7d206e10da11/go.mod h1:b+2E5dAYhXwXZwtnZ6UAqBI28+e2cm9otk0dWdXHAEo=
golang.org/x/tools v0.38.0 h1:Hx2Xv8hISq8Lm16jvBZ2VQf+RLmbd7wVUsALibYI/IQ=
golang.org/x/tools v0.38.0/go.mod h1:yEsQ/d/YK8cjh0L6rZlY8tgtlKiBNTL14pGDJPJpYQs=
golang.org/x/tools v0.40.0 h1:yLkxfA+Qnul4cs9QA3KnlFu0lVmd8JJfoq+E41uSutA=
golang.org/x/tools v0.40.0/go.mod h1:Ik/tzLRlbscWpqqMRjyWYDisX8bG13FrdXp3o4Sr9lc=
golang.org/x/xerrors v0.0.0-20190717185122-a985d3407aa7/go.mod h1:I/5z698sn9Ka8TeJc9MKroUUfqBBauWjQqLJ2OPfmY0=
golang.org/x/xerrors v0.0.0-20191204190536-9bdfabe68543/go.mod h1:I/5z698sn9Ka8TeJc9MKroUUfqBBauWjQqLJ2OPfmY0=
gonum.org/v1/gonum v0.16.0 h1:5+ul4Swaf3ESvrOnidPp4GZbzf0mxVQpDCYUQE7OJfk=
Expand Down
92 changes: 92 additions & 0 deletions internal/codemode/executor.go
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -4,11 +4,17 @@ import (
"bytes"
"context"
"fmt"
"go/parser"
"go/token"
"os"
"os/exec"
"path/filepath"
"sort"
"strings"
"syscall"
"time"

"golang.org/x/tools/imports"
)

// mcpSDKVersion is the version of the MCP SDK to use in generated go.mod
Expand Down Expand Up @@ -60,11 +66,15 @@ func ExecuteCode(ctx context.Context, servers []ServerToolsInfo, llmCode string,
return ExecutionResult{}, fmt.Errorf("writing run.go: %w", err)
}

// Auto-correct imports
importNote := autoCorrectImports(ctx, tempDir, "run.go")

// Run go mod tidy
tidyResult, err := runCommand(ctx, tempDir, "go", "mod", "tidy")
if err != nil {
return ExecutionResult{}, fmt.Errorf("running go mod tidy: %w", err)
}
tidyResult.Output += importNote
if tidyResult.ExitCode != 0 {
return ExecutionResult{
Output: tidyResult.Output,
Expand All @@ -78,6 +88,7 @@ func ExecuteCode(ctx context.Context, servers []ServerToolsInfo, llmCode string,
if err != nil {
return ExecutionResult{}, fmt.Errorf("running go build: %w", err)
}
buildResult.Output += importNote
if buildResult.ExitCode != 0 {
return ExecutionResult{
Output: buildResult.Output,
Expand All @@ -87,6 +98,7 @@ func ExecuteCode(ctx context.Context, servers []ServerToolsInfo, llmCode string,

// Execute the built binary with timeout and graceful shutdown
result, err := runProgramWithTimeout(ctx, binaryPath, timeoutSecs)
result.Output += importNote
if err != nil {
return result, err
}
Expand Down Expand Up @@ -190,3 +202,83 @@ func classifyExitCode(result ExecutionResult) error {
return RecoverableError{Output: result.Output, ExitCode: result.ExitCode}
}
}

// autoCorrectImports runs goimports (via golang.org/x/tools/imports) on the file
// and returns a notification message listing added/removed packages.
func autoCorrectImports(ctx context.Context, dir, filename string) string {
filePath := filepath.Join(dir, filename)
orig, err := os.ReadFile(filePath)
if err != nil {
return ""
}

origImports := extractImports(orig)

// Process the file using golang.org/x/tools/imports
newContent, err := imports.Process(filePath, orig, nil)
if err != nil {
// If processing fails (e.g. syntax errors), let the compiler catch it
return ""
}

if bytes.Equal(orig, newContent) {
return ""
}

if err := os.WriteFile(filePath, newContent, 0644); err != nil {
return ""
}

newImports := extractImports(newContent)
return formatImportChanges(filename, origImports, newImports)
}

// extractImports parses Go source and returns a set of import paths.
func extractImports(src []byte) map[string]bool {
fset := token.NewFileSet()
f, err := parser.ParseFile(fset, "", src, parser.ImportsOnly)
if err != nil {
return nil
}

imps := make(map[string]bool)
for _, imp := range f.Imports {
// Remove quotes from import path
path := strings.Trim(imp.Path.Value, "`\"`")
imps[path] = true
}
return imps
}

// formatImportChanges generates a message describing which imports were added/removed.
func formatImportChanges(filename string, origImports, newImports map[string]bool) string {
var added, removed []string

for pkg := range newImports {
if !origImports[pkg] {
added = append(added, pkg)
}
}
for pkg := range origImports {
if !newImports[pkg] {
removed = append(removed, pkg)
}
}

if len(added) == 0 && len(removed) == 0 {
return ""
}

sort.Strings(added)
sort.Strings(removed)

var sb strings.Builder
sb.WriteString(fmt.Sprintf("\n\nNote: Imports in %s were auto-corrected.", filename))
if len(added) > 0 {
sb.WriteString(fmt.Sprintf("\n Added: %s", strings.Join(added, ", ")))
}
if len(removed) > 0 {
sb.WriteString(fmt.Sprintf("\n Removed: %s", strings.Join(removed, ", ")))
}
return sb.String()
}
39 changes: 39 additions & 0 deletions internal/codemode/executor_test.go
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -523,3 +523,42 @@ func Run(ctx context.Context) error {
t.Fatalf("go build error: %v\n%s\n\nGenerated main.go:\n%s", err, out, mainGo)
}
}


func TestExecuteCode_AutoCorrectImports(t *testing.T) {
// Only run this test if goimports is installed
ctx := context.Background()
llmCode := `package main

import (
"context"
)

func Run(ctx context.Context) error {
// fmt is missing, but goimports should add it
fmt.Println("Imports corrected")
return nil
}
`

result, err := ExecuteCode(ctx, nil, llmCode, 30)
if err != nil {
t.Fatalf("ExecuteCode() error: %v, output: %s", err, result.Output)
}

if result.ExitCode != 0 {
t.Errorf("ExitCode = %d, want 0; output: %s", result.ExitCode, result.Output)
}

if !strings.Contains(result.Output, "Imports in run.go were auto-corrected") {
t.Errorf("Output = %q, want to contain auto-correction note", result.Output)
}

if !strings.Contains(result.Output, "Added: fmt") {
t.Errorf("Output = %q, want to contain 'Added: fmt'", result.Output)
}

if !strings.Contains(result.Output, "Imports corrected") {
t.Errorf("Output = %q, want to contain program output 'Imports corrected'", result.Output)
}
}
10 changes: 6 additions & 4 deletions internal/codemode/tool_test.go
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -52,19 +52,21 @@ func Run(ctx context.Context) error {
},
wantOutputSub: "syntax error",
},
{
name: "compilation error - missing import",
{
name: "compilation error - undefined variable",
input: ExecuteGoCodeInput{
Code: `package main

import "context"

func Run(ctx context.Context) error {
fmt.Println("missing import")
_ = bar // undefined
return nil
}
`,
ExecutionTimeout: 30,
},
wantOutputSub: "undefined: fmt",
wantOutputSub: "undefined: bar",
},
{
name: "Run returns error - exit code 1",
Expand Down