From 22823f3f9d544d84bd61cea9c862a54b3b8c146e Mon Sep 17 00:00:00 2001 From: Yumiue <229866007@qq.com> Date: Wed, 1 Apr 2026 22:27:01 +0800 Subject: [PATCH 1/3] feat(context): inject project rules and system state --- .gitattributes | 12 +- docs/neocode-coding-agent-mvp-architecture.md | 15 +- docs/runtime-provider-event-flow.md | 30 +++ internal/context/builder.go | 23 ++- internal/context/builder_test.go | 21 ++- internal/context/metadata.go | 25 +++ internal/context/prompt.go | 14 ++ internal/context/source_rules.go | 178 ++++++++++++++++++ internal/context/source_rules_test.go | 121 ++++++++++++ internal/context/source_system.go | 82 ++++++++ internal/context/source_system_test.go | 70 +++++++ internal/context/types.go | 2 +- internal/runtime/runtime.go | 7 +- internal/runtime/runtime_test.go | 13 +- 14 files changed, 598 insertions(+), 15 deletions(-) create mode 100644 internal/context/metadata.go create mode 100644 internal/context/source_rules.go create mode 100644 internal/context/source_rules_test.go create mode 100644 internal/context/source_system.go create mode 100644 internal/context/source_system_test.go diff --git a/.gitattributes b/.gitattributes index f8601b71..2e502d2a 100644 --- a/.gitattributes +++ b/.gitattributes @@ -1,2 +1,12 @@ +# 统一仓库内常见文本文件为 LF，避免 Windows 下出现无意义的换行警告 +*.go text eol=lf +*.md text eol=lf +*.yml text eol=lf +*.yaml text eol=lf +go.mod text eol=lf +go.sum text eol=lf +.gitignore text eol=lf +.gitattributes text eol=lf + # 强制 scripts 目录下的 Shell 脚本始终使用 LF 换行符 -scripts/*.sh text eol=lf \ No newline at end of file +scripts/*.sh text eol=lf diff --git a/docs/neocode-coding-agent-mvp-architecture.md b/docs/neocode-coding-agent-mvp-architecture.md index a10e07a4..4f394554 100644 --- a/docs/neocode-coding-agent-mvp-architecture.md +++ b/docs/neocode-coding-agent-mvp-architecture.md @@ -6,13 +6,14 @@ `用户输入 -> Agent 推理 -> 调用工具 -> 获取结果 -> 继续推理 -> UI 展示` -MVP 聚焦五个模块： +MVP 聚焦六个模块： 1. provider：统一不同模型/API 的调用方式 2. TUI：用户交互入口，承载输入、对话、侧边栏、会话 3. tools：统一工具定义、参数校验、执行与结果封装 4. config：管理本地配置、provider 切换、模型选择 -5. agent runtime：驱动整个 agent loop，是系统核心 +5. context：负责 system prompt、显式上下文源与历史消息裁剪 +6. agent runtime：驱动整个 agent loop，是系统核心 --- @@ -46,6 +47,7 @@ flowchart LR - TUI：负责交互和渲染 - Application：负责启动和依赖注入 - Runtime：负责 Agent Loop 和状态编排 +- Context：负责模型请求前的上下文构建 - Provider：负责模型调用抽象 - Tool Manager：负责工具注册、校验、执行 - Config：负责配置加载与选择 @@ -295,7 +297,7 @@ type UserInput struct { Runtime 内部建议拆分： - `SessionStore`：管理会话和消息历史 -- `PromptBuilder`：组装 prompt/messages/tools +- `context.Builder`：组装核心 prompt、显式上下文源与裁剪后的消息 - `Executor`：执行 loop - `EventBus`：向 TUI 推送运行事件 @@ -379,6 +381,13 @@ sequenceDiagram │ │ ├── loader.go │ │ ├── model.go │ │ └── validate.go +│ ├── context/ +│ │ ├── builder.go +│ │ ├── metadata.go +│ │ ├── prompt.go +│ │ ├── source_rules.go +│ │ ├── source_system.go +│ │ └── trim.go │ ├── provider/ │ │ ├── provider.go │ │ ├── openai/ diff --git a/docs/runtime-provider-event-flow.md b/docs/runtime-provider-event-flow.md index e4f0a982..c39c740a 100644 --- a/docs/runtime-provider-event-flow.md +++ b/docs/runtime-provider-event-flow.md @@ -22,6 +22,36 @@ 8. 执行返回的工具调用，并保存每一个工具结果。 9. 如果仍需继续推理，则进入下一轮；否则结束。 +### Context Builder 输入与职责 + +- `runtime` 只向 `context.Builder` 传递本轮所需元数据： + - 历史消息 + - `workdir` + - `shell` + - 当前 `provider` + - 当前 `model` +- `context.Builder` 负责统一组装： + - 固定核心 system prompt + - 从 `workdir` 向上发现的 `AGENTS.md` + - 系统状态摘要（`workdir` / `shell` / `provider` / `model` / git branch / git dirty） + - 裁剪后的历史消息 +- `runtime` 不直接读取规则文件，也不直接查询 git 状态。 +- `provider` 只消费最终生成的 `SystemPrompt`、消息列表和工具 schema，不感知上下文来源。 + +### System Prompt 注入顺序 + +当前 `system prompt` 按以下顺序拼装： + +1. 固定核心指令 +2. `Project Rules` section +3. `System State` section + +其中： + +- 规则文件只支持大写文件名 `AGENTS.md` +- 多份命中结果按“从全局到局部”的顺序注入 +- git 只注入摘要，不注入完整 `git status` + ## 流式桥接 - Provider 发出 `StreamEvent` diff --git a/internal/context/builder.go b/internal/context/builder.go index de106976..b87777f5 100644 --- a/internal/context/builder.go +++ b/internal/context/builder.go @@ -3,11 +3,15 @@ package context import "context" // DefaultBuilder preserves the current runtime context-building behavior. -type DefaultBuilder struct{} +type DefaultBuilder struct { + gitRunner gitCommandRunner +} // NewBuilder returns the default context builder implementation. func NewBuilder() Builder { - return &DefaultBuilder{} + return &DefaultBuilder{ + gitRunner: runGitCommand, + } } // Build assembles the provider-facing context for the current round. @@ -16,8 +20,19 @@ func (b *DefaultBuilder) Build(ctx context.Context, input BuildInput) (BuildResu return BuildResult{}, err } + rules, err := loadProjectRules(ctx, input.Metadata.Workdir) + if err != nil { + return BuildResult{}, err + } + + systemState := collectSystemState(ctx, input.Metadata, b.gitRunner) + return BuildResult{ - SystemPrompt: defaultSystemPrompt(), - Messages: trimMessages(input.Messages), + SystemPrompt: composeSystemPrompt( + defaultSystemPrompt(), + renderProjectRulesSection(rules), + renderSystemStateSection(systemState), + ), + Messages: trimMessages(input.Messages), }, nil } diff --git a/internal/context/builder_test.go b/internal/context/builder_test.go index fe23ed26..a7b80995 100644 --- a/internal/context/builder_test.go +++ b/internal/context/builder_test.go @@ -3,6 +3,7 @@ package context import ( stdcontext "context" "fmt" + "strings" "testing" "neo-code/internal/provider" @@ -16,7 +17,12 @@ func TestDefaultBuilderBuild(t *testing.T) { Messages: []provider.Message{ {Role: "user", Content: "hello"}, }, - Workdir: t.TempDir(), + Metadata: Metadata{ + Workdir: t.TempDir(), + Shell: "powershell", + Provider: "openai", + Model: "gpt-5.4", + }, } got, err := builder.Build(stdcontext.Background(), input) @@ -26,8 +32,17 @@ func TestDefaultBuilderBuild(t *testing.T) { if got.SystemPrompt == "" { t.Fatalf("expected non-empty system prompt") } - if got.SystemPrompt != defaultSystemPrompt() { - t.Fatalf("expected default prompt to remain unchanged") + if !strings.Contains(got.SystemPrompt, defaultSystemPrompt()) { + t.Fatalf("expected default prompt to remain in composed prompt") + } + if !strings.Contains(got.SystemPrompt, "## System State") { + t.Fatalf("expected system state section in composed prompt") + } + if strings.Contains(got.SystemPrompt, "## Project Rules") { + t.Fatalf("did not expect project rules section without AGENTS.md") + } + if !strings.Contains(got.SystemPrompt, input.Metadata.Workdir) { + t.Fatalf("expected workdir in system state section") } if len(got.Messages) != 1 { t.Fatalf("expected 1 message, got %d", len(got.Messages)) diff --git a/internal/context/metadata.go b/internal/context/metadata.go new file mode 100644 index 00000000..bec1f559 --- /dev/null +++ b/internal/context/metadata.go @@ -0,0 +1,25 @@ +package context + +// Metadata contains the non-message runtime state needed by context sources. +type Metadata struct { + Workdir string + Shell string + Provider string + Model string +} + +// GitState is the summarized git metadata exposed to the prompt builder. +type GitState struct { + Available bool + Branch string + Dirty bool +} + +// SystemState is the summarized runtime metadata exposed to the prompt builder. +type SystemState struct { + Workdir string + Shell string + Provider string + Model string + Git GitState +} diff --git a/internal/context/prompt.go b/internal/context/prompt.go index db01563c..c3c44e2f 100644 --- a/internal/context/prompt.go +++ b/internal/context/prompt.go @@ -1,5 +1,7 @@ package context +import "strings" + func defaultSystemPrompt() string { return `You are NeoCode, a local coding agent. @@ -8,3 +10,15 @@ func defaultSystemPrompt() string { When a tool fails, inspect the error and continue safely. Stay within the workspace and avoid destructive behavior unless clearly requested.` } + +func composeSystemPrompt(parts ...string) string { + sections := make([]string, 0, len(parts)) + for _, part := range parts { + part = strings.TrimSpace(part) + if part == "" { + continue + } + sections = append(sections, part) + } + return strings.Join(sections, "\n\n") +} diff --git a/internal/context/source_rules.go b/internal/context/source_rules.go new file mode 100644 index 00000000..621a999a --- /dev/null +++ b/internal/context/source_rules.go @@ -0,0 +1,178 @@ +package context + +import ( + "context" + "fmt" + "os" + "path/filepath" + "strings" + "unicode/utf8" +) + +const ( + ruleFileName = "AGENTS.md" + maxRuleFileRunes = 4000 + maxTotalRuleRunes = 12000 +) + +type ruleDocument struct { + Path string + Content string + Truncated bool +} + +func loadProjectRules(ctx context.Context, workdir string) ([]ruleDocument, error) { + paths, err := discoverRuleFiles(ctx, workdir) + if err != nil { + return nil, err + } + + return loadRuleDocuments(ctx, paths, os.ReadFile) +} + +func loadRuleDocuments(ctx context.Context, paths []string, readFile func(string) ([]byte, error)) ([]ruleDocument, error) { + documents := make([]ruleDocument, 0, len(paths)) + for _, path := range paths { + if err := ctx.Err(); err != nil { + return nil, err + } + + data, err := readFile(path) + if err != nil { + return nil, fmt.Errorf("context: read %s: %w", path, err) + } + + content, truncated := truncateRunes(strings.TrimSpace(string(data)), maxRuleFileRunes) + documents = append(documents, ruleDocument{ + Path: path, + Content: content, + Truncated: truncated, + }) + } + + return documents, nil +} + +func discoverRuleFiles(ctx context.Context, workdir string) ([]string, error) { + workdir = strings.TrimSpace(workdir) + if workdir == "" { + return nil, nil + } + + dir := filepath.Clean(workdir) + if info, err := os.Stat(dir); err == nil && !info.IsDir() { + dir = filepath.Dir(dir) + } + + paths := make([]string, 0, 4) + for { + if err := ctx.Err(); err != nil { + return nil, err + } + + match, err := findExactRuleFile(dir) + if err != nil { + return nil, err + } + if match != "" { + paths = append(paths, match) + } + + parent := filepath.Dir(dir) + if parent == dir { + break + } + dir = parent + } + + for i, j := 0, len(paths)-1; i < j; i, j = i+1, j-1 { + paths[i], paths[j] = paths[j], paths[i] + } + + return paths, nil +} + +func findExactRuleFile(dir string) (string, error) { + entries, err := os.ReadDir(dir) + if err != nil { + if os.IsNotExist(err) { + return "", nil + } + return "", fmt.Errorf("context: read dir %s: %w", dir, err) + } + + for _, entry := range entries { + if entry.IsDir() { + continue + } + if entry.Name() == ruleFileName { + return filepath.Join(dir, entry.Name()), nil + } + } + + return "", nil +} + +func renderProjectRulesSection(documents []ruleDocument) string { + if len(documents) == 0 { + return "" + } + + var builder strings.Builder + builder.WriteString("## Project Rules\n") + + remaining := maxTotalRuleRunes + totalTruncated := false + for _, document := range documents { + if remaining <= 0 { + totalTruncated = true + break + } + + content := document.Content + if runeCount(content) > remaining { + content, _ = truncateRunes(content, remaining) + totalTruncated = true + } + remaining -= runeCount(content) + + builder.WriteString("\n### ") + builder.WriteString(document.Path) + builder.WriteString("\n") + if content != "" { + builder.WriteString("\n") + builder.WriteString(content) + builder.WriteString("\n") + } + if document.Truncated { + builder.WriteString("\n[truncated to fit per-file limit]\n") + } + + if content != document.Content { + totalTruncated = true + break + } + } + + if totalTruncated { + builder.WriteString("\n[additional project rules truncated to fit total limit]\n") + } + + return strings.TrimSpace(builder.String()) +} + +func truncateRunes(input string, max int) (string, bool) { + if max <= 0 { + return "", input != "" + } + if runeCount(input) <= max { + return input, false + } + + runes := []rune(input) + return string(runes[:max]), true +} + +func runeCount(input string) int { + return utf8.RuneCountInString(input) +} diff --git a/internal/context/source_rules_test.go b/internal/context/source_rules_test.go new file mode 100644 index 00000000..49bbcc21 --- /dev/null +++ b/internal/context/source_rules_test.go @@ -0,0 +1,121 @@ +package context + +import ( + "context" + "errors" + "os" + "path/filepath" + "strings" + "testing" +) + +func TestLoadProjectRulesOrdersGlobalToLocal(t *testing.T) { + t.Parallel() + + root := t.TempDir() + nested := filepath.Join(root, "a", "b") + if err := os.MkdirAll(nested, 0o755); err != nil { + t.Fatalf("mkdir nested: %v", err) + } + + rootRules := filepath.Join(root, ruleFileName) + localRules := filepath.Join(root, "a", ruleFileName) + if err := os.WriteFile(rootRules, []byte("root-rules"), 0o644); err != nil { + t.Fatalf("write root rules: %v", err) + } + if err := os.WriteFile(localRules, []byte("local-rules"), 0o644); err != nil { + t.Fatalf("write local rules: %v", err) + } + + documents, err := loadProjectRules(context.Background(), nested) + if err != nil { + t.Fatalf("loadProjectRules() error = %v", err) + } + if len(documents) != 2 { + t.Fatalf("expected 2 rule documents, got %d", len(documents)) + } + if documents[0].Path != rootRules || documents[1].Path != localRules { + t.Fatalf("expected global-to-local order, got %+v", documents) + } + + section := renderProjectRulesSection(documents) + rootIndex := strings.Index(section, rootRules) + localIndex := strings.Index(section, localRules) + if rootIndex < 0 || localIndex < 0 || rootIndex >= localIndex { + t.Fatalf("expected rendered rules to stay global-to-local, got %q", section) + } +} + +func TestLoadProjectRulesOnlyMatchesUppercase(t *testing.T) { + t.Parallel() + + root := t.TempDir() + nested := filepath.Join(root, "child") + if err := os.MkdirAll(nested, 0o755); err != nil { + t.Fatalf("mkdir child: %v", err) + } + + if err := os.WriteFile(filepath.Join(root, "agents.md"), []byte("wrong-case"), 0o644); err != nil { + t.Fatalf("write lowercase rules: %v", err) + } + if err := os.WriteFile(filepath.Join(nested, ruleFileName), []byte("right-case"), 0o644); err != nil { + t.Fatalf("write uppercase rules: %v", err) + } + + documents, err := loadProjectRules(context.Background(), nested) + if err != nil { + t.Fatalf("loadProjectRules() error = %v", err) + } + if len(documents) != 1 { + t.Fatalf("expected only uppercase AGENTS.md to be loaded, got %+v", documents) + } + if filepath.Base(documents[0].Path) != ruleFileName { + t.Fatalf("expected uppercase AGENTS.md match, got %q", documents[0].Path) + } + if strings.Contains(documents[0].Content, "wrong-case") { + t.Fatalf("did not expect lowercase agents.md content to be loaded") + } +} + +func TestLoadRuleDocumentsReturnsReadError(t *testing.T) { + t.Parallel() + + root := t.TempDir() + path := filepath.Join(root, ruleFileName) + if err := os.WriteFile(path, []byte("rules"), 0o644); err != nil { + t.Fatalf("write rules: %v", err) + } + + _, err := loadRuleDocuments(context.Background(), []string{path}, func(string) ([]byte, error) { + return nil, errors.New("boom") + }) + if err == nil || !strings.Contains(err.Error(), "boom") { + t.Fatalf("expected read error, got %v", err) + } +} + +func TestRenderProjectRulesSectionTruncatesSingleFileAndTotalBudget(t *testing.T) { + t.Parallel() + + largeSingle := strings.Repeat("a", maxRuleFileRunes+32) + largeTotalA := strings.Repeat("b", 7000) + largeTotalB := strings.Repeat("c", 7000) + + section := renderProjectRulesSection([]ruleDocument{ + {Path: "/repo/AGENTS.md", Content: largeSingle[:maxRuleFileRunes], Truncated: true}, + }) + if !strings.Contains(section, "[truncated to fit per-file limit]") { + t.Fatalf("expected per-file truncation marker, got %q", section) + } + + totalSection := renderProjectRulesSection([]ruleDocument{ + {Path: "/repo/root/AGENTS.md", Content: largeTotalA}, + {Path: "/repo/root/app/AGENTS.md", Content: largeTotalB}, + }) + if !strings.Contains(totalSection, "[additional project rules truncated to fit total limit]") { + t.Fatalf("expected total truncation marker, got %q", totalSection) + } + if strings.Contains(totalSection, strings.Repeat("c", 6500)) { + t.Fatalf("expected total rules section to be truncated") + } +} diff --git a/internal/context/source_system.go b/internal/context/source_system.go new file mode 100644 index 00000000..1360bbf8 --- /dev/null +++ b/internal/context/source_system.go @@ -0,0 +1,82 @@ +package context + +import ( + "context" + "fmt" + "os/exec" + "strings" +) + +type gitCommandRunner func(ctx context.Context, workdir string, args ...string) (string, error) + +func collectSystemState(ctx context.Context, metadata Metadata, runner gitCommandRunner) SystemState { + state := SystemState{ + Workdir: strings.TrimSpace(metadata.Workdir), + Shell: strings.TrimSpace(metadata.Shell), + Provider: strings.TrimSpace(metadata.Provider), + Model: strings.TrimSpace(metadata.Model), + } + + if err := ctx.Err(); err != nil { + return state + } + if runner == nil || state.Workdir == "" { + return state + } + + branch, err := runner(ctx, state.Workdir, "rev-parse", "--abbrev-ref", "HEAD") + if err != nil { + return state + } + dirty, err := runner(ctx, state.Workdir, "status", "--porcelain") + if err != nil { + return state + } + + state.Git = GitState{ + Available: true, + Branch: strings.TrimSpace(branch), + Dirty: strings.TrimSpace(dirty) != "", + } + return state +} + +func renderSystemStateSection(state SystemState) string { + lines := []string{ + "## System State", + "", + fmt.Sprintf("- workdir: `%s`", promptValue(state.Workdir)), + fmt.Sprintf("- shell: `%s`", promptValue(state.Shell)), + fmt.Sprintf("- provider: `%s`", promptValue(state.Provider)), + fmt.Sprintf("- model: `%s`", promptValue(state.Model)), + } + + if state.Git.Available { + dirty := "clean" + if state.Git.Dirty { + dirty = "dirty" + } + lines = append(lines, fmt.Sprintf("- git: branch=`%s`, dirty=`%s`", promptValue(state.Git.Branch), dirty)) + } else { + lines = append(lines, "- git: unavailable") + } + + return strings.Join(lines, "\n") +} + +func runGitCommand(ctx context.Context, workdir string, args ...string) (string, error) { + command := exec.CommandContext(ctx, "git", append([]string{"-C", workdir}, args...)...) + output, err := command.Output() + if err != nil { + return "", err + } + return string(output), nil +} + +func promptValue(value string) string { + value = strings.TrimSpace(value) + if value == "" { + return "unknown" + } + return value +} diff --git a/internal/context/source_system_test.go b/internal/context/source_system_test.go new file mode 100644 index 00000000..19515878 --- /dev/null +++ b/internal/context/source_system_test.go @@ -0,0 +1,70 @@ +package context + +import ( + "context" + "errors" + "strings" + "testing" +) + +func TestCollectSystemStateHandlesGitUnavailable(t *testing.T) { + t.Parallel() + + state := collectSystemState(context.Background(), Metadata{ + Workdir: "/workspace", + Shell: "bash", + Provider: "openai", + Model: "gpt-5.4", + }, func(ctx context.Context, workdir string, args ...string) (string, error) { + return "", errors.New("git unavailable") + }) + + if state.Git.Available { + t.Fatalf("expected git to be unavailable") + } + + section := renderSystemStateSection(state) + if !strings.Contains(section, "- git: unavailable") { + t.Fatalf("expected unavailable git section, got %q", section) + } +} + +func TestCollectSystemStateIncludesGitSummary(t *testing.T) { + t.Parallel() + + runner := func(ctx context.Context, workdir string, args ...string) (string, error) { + switch strings.Join(args, " ") { + case "rev-parse --abbrev-ref HEAD": + return "feature/context\n", nil + case "status --porcelain": + return " M internal/context/builder.go\n", nil + default: + return "", errors.New("unexpected git command") + } + } + + state := collectSystemState(context.Background(), Metadata{ + Workdir: "/workspace", + Shell: "bash", + Provider: "openai", + Model: "gpt-5.4", + }, runner) + + if !state.Git.Available { + t.Fatalf("expected git to be available") + } + if state.Git.Branch != "feature/context" { + t.Fatalf("expected branch to be trimmed, got %q", state.Git.Branch) + } + if !state.Git.Dirty { + t.Fatalf("expected dirty git state") + } + + section := renderSystemStateSection(state) + if !strings.Contains(section, "branch=`feature/context`") { + t.Fatalf("expected branch in system section, got %q", section) + } + if !strings.Contains(section, "dirty=`dirty`") { + t.Fatalf("expected dirty marker in system section, got %q", section) + } +} diff --git a/internal/context/types.go b/internal/context/types.go index 0c1e7a96..6eb516dd 100644 --- a/internal/context/types.go +++ b/internal/context/types.go @@ -14,7 +14,7 @@ type Builder interface { // BuildInput contains the runtime state needed to assemble model context. type BuildInput struct { Messages []provider.Message - Workdir string + Metadata Metadata } // BuildResult is the provider-facing context produced for a single round. diff --git a/internal/runtime/runtime.go b/internal/runtime/runtime.go index 6e0da509..49825ac1 100644 --- a/internal/runtime/runtime.go +++ b/internal/runtime/runtime.go @@ -130,7 +130,12 @@ func (s *Service) Run(ctx context.Context, input UserInput) error { builtContext, err := s.contextBuilder.Build(ctx, agentcontext.BuildInput{ Messages: session.Messages, - Workdir: cfg.Workdir, + Metadata: agentcontext.Metadata{ + Workdir: cfg.Workdir, + Shell: cfg.Shell, + Provider: cfg.SelectedProvider, + Model: cfg.CurrentModel, + }, }) if err != nil { return s.handleRunError(ctx, input.RunID, session.ID, err) diff --git a/internal/runtime/runtime_test.go b/internal/runtime/runtime_test.go index 4f3eb869..3bff01ae 100644 --- a/internal/runtime/runtime_test.go +++ b/internal/runtime/runtime_test.go @@ -413,8 +413,17 @@ func TestServiceRunDelegatesToContextBuilder(t *testing.T) { if builder.callCount != 1 { t.Fatalf("expected builder to be called once, got %d", builder.callCount) } - if builder.lastInput.Workdir == "" { - t.Fatalf("expected workdir to be forwarded to builder") + if builder.lastInput.Metadata.Workdir == "" { + t.Fatalf("expected workdir to be forwarded to builder metadata") + } + if builder.lastInput.Metadata.Shell == "" { + t.Fatalf("expected shell to be forwarded to builder metadata") + } + if builder.lastInput.Metadata.Provider == "" { + t.Fatalf("expected provider to be forwarded to builder metadata") + } + if builder.lastInput.Metadata.Model == "" { + t.Fatalf("expected model to be forwarded to builder metadata") } if len(builder.lastInput.Messages) != 1 || builder.lastInput.Messages[0].Content != "hello" { t.Fatalf("expected persisted session messages to be forwarded, got %+v", builder.lastInput.Messages) From d89800128bb418ed3ebc8016ec93b1648b1cb661 Mon Sep 17 00:00:00 2001 From: Yumiue <229866007@qq.com> Date: Wed, 1 Apr 2026 22:45:49 +0800 Subject: [PATCH 2/3] fix(context): harden explicit context sources --- internal/context/builder.go | 5 +- internal/context/source_rules.go | 109 +++++++++++++++++++------ internal/context/source_rules_test.go | 43 ++++++++++ internal/context/source_system.go | 23 ++++-- internal/context/source_system_test.go | 26 +++++- 5 files changed, 173 insertions(+), 33 deletions(-) diff --git a/internal/context/builder.go b/internal/context/builder.go index b87777f5..bc57052b 100644 --- a/internal/context/builder.go +++ b/internal/context/builder.go @@ -25,7 +25,10 @@ func (b *DefaultBuilder) Build(ctx context.Context, input BuildInput) (BuildResu return BuildResult{}, err } - systemState := collectSystemState(ctx, input.Metadata, b.gitRunner) + systemState, err := collectSystemState(ctx, input.Metadata, b.gitRunner) + if err != nil { + return BuildResult{}, err + } return BuildResult{ SystemPrompt: composeSystemPrompt( diff --git a/internal/context/source_rules.go b/internal/context/source_rules.go index 621a999a..a666132b 100644 --- a/internal/context/source_rules.go +++ b/internal/context/source_rules.go @@ -21,6 +21,8 @@ type ruleDocument struct { Truncated bool } +type ruleFileFinder func(string) (string, error) + func loadProjectRules(ctx context.Context, workdir string) ([]ruleDocument, error) { paths, err := discoverRuleFiles(ctx, workdir) if err != nil { @@ -54,6 +56,10 @@ func loadRuleDocuments(ctx context.Context, paths []string, readFile func(string } func discoverRuleFiles(ctx context.Context, workdir string) ([]string, error) { + return discoverRuleFilesWithFinder(ctx, workdir, findExactRuleFile) +} + +func discoverRuleFilesWithFinder(ctx context.Context, workdir string, finder ruleFileFinder) ([]string, error) { workdir = strings.TrimSpace(workdir) if workdir == "" { return nil, nil @@ -70,9 +76,9 @@ func discoverRuleFiles(ctx context.Context, workdir string) ([]string, error) { return nil, err } - match, err := findExactRuleFile(dir) + match, err := finder(dir) if err != nil { - return nil, err + break } if match != "" { paths = append(paths, match) @@ -118,47 +124,102 @@ func renderProjectRulesSection(documents []ruleDocument) string { return "" } + const totalTruncationNotice = "\n[additional project rules truncated to fit total limit]\n" + var builder strings.Builder builder.WriteString("## Project Rules\n") remaining := maxTotalRuleRunes - totalTruncated := false + totalBudgetTruncated := false for _, document := range documents { if remaining <= 0 { - totalTruncated = true + totalBudgetTruncated = true break } - content := document.Content - if runeCount(content) > remaining { - content, _ = truncateRunes(content, remaining) - totalTruncated = true + fullChunk := renderRuleDocumentChunk(document) + fullChunkRunes := runeCount(fullChunk) + if fullChunkRunes <= remaining { + builder.WriteString(fullChunk) + remaining -= fullChunkRunes + continue } - remaining -= runeCount(content) - builder.WriteString("\n### ") - builder.WriteString(document.Path) - builder.WriteString("\n") - if content != "" { - builder.WriteString("\n") - builder.WriteString(content) - builder.WriteString("\n") + totalBudgetTruncated = true + chunkBudget := remaining + if noticeRunes := runeCount(totalTruncationNotice); noticeRunes < chunkBudget { + chunkBudget -= noticeRunes } - if document.Truncated { - builder.WriteString("\n[truncated to fit per-file limit]\n") + chunk := renderRuleDocumentChunkWithinBudget(document, chunkBudget) + builder.WriteString(chunk) + remaining -= runeCount(chunk) + break + } + + if totalBudgetTruncated { + if runeCount(totalTruncationNotice) <= remaining { + builder.WriteString(totalTruncationNotice) } + } - if content != document.Content { - totalTruncated = true - break + return strings.TrimSpace(builder.String()) +} + +func renderRuleDocumentChunk(document ruleDocument) string { + var builder strings.Builder + builder.WriteString("\n### ") + builder.WriteString(document.Path) + builder.WriteString("\n") + if document.Content != "" { + builder.WriteString("\n") + builder.WriteString(document.Content) + builder.WriteString("\n") + } + if document.Truncated { + builder.WriteString("\n[truncated to fit per-file limit]\n") + } + + return builder.String() +} + +func renderRuleDocumentChunkWithinBudget(document ruleDocument, budget int) string { + if budget <= 0 { + return "" + } + + header := "\n### " + document.Path + "\n" + headerRunes := runeCount(header) + if headerRunes > budget { + return "" + } + + bodyBudget := budget - headerRunes + content := document.Content + if runeCount(content) > bodyBudget { + content, _ = truncateRunes(content, bodyBudget) + } + + var body strings.Builder + if content != "" { + body.WriteString("\n") + body.WriteString(content) + body.WriteString("\n") + } + if document.Truncated { + perFileNotice := "\n[truncated to fit per-file limit]\n" + if runeCount(body.String())+runeCount(perFileNotice) <= bodyBudget { + body.WriteString(perFileNotice) } } - if totalTruncated { - builder.WriteString("\n[additional project rules truncated to fit total limit]\n") + bodyRunes := runeCount(body.String()) + if bodyRunes > bodyBudget { + bodyText, _ := truncateRunes(body.String(), bodyBudget) + body.Reset() + body.WriteString(bodyText) } - return strings.TrimSpace(builder.String()) + return header + body.String() } func truncateRunes(input string, max int) (string, bool) { diff --git a/internal/context/source_rules_test.go b/internal/context/source_rules_test.go index 49bbcc21..12cbda99 100644 --- a/internal/context/source_rules_test.go +++ b/internal/context/source_rules_test.go @@ -94,6 +94,45 @@ func TestLoadRuleDocumentsReturnsReadError(t *testing.T) { } } +func TestDiscoverRuleFilesStopsOnDirectoryReadError(t *testing.T) { + t.Parallel() + + root := t.TempDir() + nested := filepath.Join(root, "a", "b") + if err := os.MkdirAll(nested, 0o755); err != nil { + t.Fatalf("mkdir nested: %v", err) + } + + rootRules := filepath.Join(root, ruleFileName) + localRules := filepath.Join(root, "a", ruleFileName) + if err := os.WriteFile(rootRules, []byte("root-rules"), 0o644); err != nil { + t.Fatalf("write root rules: %v", err) + } + if err := os.WriteFile(localRules, []byte("local-rules"), 0o644); err != nil { + t.Fatalf("write local rules: %v", err) + } + + permissionErr := errors.New("permission denied") + paths, err := discoverRuleFilesWithFinder(context.Background(), nested, func(dir string) (string, error) { + switch dir { + case nested: + return "", nil + case filepath.Join(root, "a"): + return localRules, nil + case root: + return "", permissionErr + default: + return "", nil + } + }) + if err != nil { + t.Fatalf("discoverRuleFilesWithFinder() error = %v", err) + } + if len(paths) != 1 || paths[0] != localRules { + t.Fatalf("expected traversal to stop after read error and keep collected paths, got %+v", paths) + } +} + func TestRenderProjectRulesSectionTruncatesSingleFileAndTotalBudget(t *testing.T) { t.Parallel() @@ -118,4 +157,8 @@ func TestRenderProjectRulesSectionTruncatesSingleFileAndTotalBudget(t *testing.T if strings.Contains(totalSection, strings.Repeat("c", 6500)) { t.Fatalf("expected total rules section to be truncated") } + body := strings.TrimPrefix(totalSection, "## Project Rules\n") + if runeCount(body) > maxTotalRuleRunes { + t.Fatalf("expected rendered rules body to respect total rune budget, got %d > %d", runeCount(body), maxTotalRuleRunes) + } } diff --git a/internal/context/source_system.go b/internal/context/source_system.go index 1360bbf8..ebc2a042 100644 --- a/internal/context/source_system.go +++ b/internal/context/source_system.go @@ -2,6 +2,7 @@ package context import ( "context" + "errors" "fmt" "os/exec" "strings" @@ -9,7 +10,7 @@ import ( type gitCommandRunner func(ctx context.Context, workdir string, args ...string) (string, error) -func collectSystemState(ctx context.Context, metadata Metadata, runner gitCommandRunner) SystemState { +func collectSystemState(ctx context.Context, metadata Metadata, runner gitCommandRunner) (SystemState, error) { state := SystemState{ Workdir: strings.TrimSpace(metadata.Workdir), Shell: strings.TrimSpace(metadata.Shell), @@ -18,19 +19,25 @@ func collectSystemState(ctx context.Context, metadata Metadata, runner gitComman } if err := ctx.Err(); err != nil { - return state + return state, err } if runner == nil || state.Workdir == "" { - return state + return state, nil } branch, err := runner(ctx, state.Workdir, "rev-parse", "--abbrev-ref", "HEAD") if err != nil { - return state + if isContextError(err) { + return state, err + } + return state, nil } dirty, err := runner(ctx, state.Workdir, "status", "--porcelain") if err != nil { - return state + if isContextError(err) { + return state, err + } + return state, nil } state.Git = GitState{ @@ -38,7 +45,7 @@ func collectSystemState(ctx context.Context, metadata Metadata, runner gitComman Branch: strings.TrimSpace(branch), Dirty: strings.TrimSpace(dirty) != "", } - return state + return state, nil } func renderSystemStateSection(state SystemState) string { @@ -80,3 +87,7 @@ func promptValue(value string) string { } return value } + +func isContextError(err error) bool { + return errors.Is(err, context.Canceled) || errors.Is(err, context.DeadlineExceeded) +} diff --git a/internal/context/source_system_test.go b/internal/context/source_system_test.go index 19515878..2a150dae 100644 --- a/internal/context/source_system_test.go +++ b/internal/context/source_system_test.go @@ -10,7 +10,7 @@ import ( func TestCollectSystemStateHandlesGitUnavailable(t *testing.T) { t.Parallel() - state := collectSystemState(context.Background(), Metadata{ + state, err := collectSystemState(context.Background(), Metadata{ Workdir: "/workspace", Shell: "bash", Provider: "openai", @@ -18,6 +18,9 @@ func TestCollectSystemStateHandlesGitUnavailable(t *testing.T) { }, func(ctx context.Context, workdir string, args ...string) (string, error) { return "", errors.New("git unavailable") }) + if err != nil { + t.Fatalf("collectSystemState() error = %v", err) + } if state.Git.Available { t.Fatalf("expected git to be unavailable") @@ -43,12 +46,15 @@ func TestCollectSystemStateIncludesGitSummary(t *testing.T) { } } - state := collectSystemState(context.Background(), Metadata{ + state, err := collectSystemState(context.Background(), Metadata{ Workdir: "/workspace", Shell: "bash", Provider: "openai", Model: "gpt-5.4", }, runner) + if err != nil { + t.Fatalf("collectSystemState() error = %v", err) + } if !state.Git.Available { t.Fatalf("expected git to be available") @@ -68,3 +74,19 @@ func TestCollectSystemStateIncludesGitSummary(t *testing.T) { t.Fatalf("expected dirty marker in system section, got %q", section) } } + +func TestCollectSystemStateReturnsContextError(t *testing.T) { + t.Parallel() + + ctx, cancel := context.WithCancel(context.Background()) + cancel() + + _, err := collectSystemState(ctx, Metadata{ + Workdir: "/workspace", + }, func(ctx context.Context, workdir string, args ...string) (string, error) { + return "", ctx.Err() + }) + if !errors.Is(err, context.Canceled) { + t.Fatalf("expected context canceled error, got %v", err) + } +} From e3042d34b6034905dd4179c0ef6f192fe4d87c38 Mon Sep 17 00:00:00 2001 From: Yumiue <229866007@qq.com> Date: Wed, 1 Apr 2026 22:54:25 +0800 Subject: [PATCH 3/3] test(context): avoid hardcoded provider metadata --- internal/context/builder_test.go | 7 +------ internal/context/source_system_test.go | 18 +++--------------- internal/context/test_helpers_test.go | 13 +++++++++++++ 3 files changed, 17 insertions(+), 21 deletions(-) create mode 100644 internal/context/test_helpers_test.go diff --git a/internal/context/builder_test.go b/internal/context/builder_test.go index a7b80995..4ca49c05 100644 --- a/internal/context/builder_test.go +++ b/internal/context/builder_test.go @@ -17,12 +17,7 @@ func TestDefaultBuilderBuild(t *testing.T) { Messages: []provider.Message{ {Role: "user", Content: "hello"}, }, - Metadata: Metadata{ - Workdir: t.TempDir(), - Shell: "powershell", - Provider: "openai", - Model: "gpt-5.4", - }, + Metadata: testMetadata(t.TempDir()), } got, err := builder.Build(stdcontext.Background(), input) diff --git a/internal/context/source_system_test.go b/internal/context/source_system_test.go index 2a150dae..df40ca38 100644 --- a/internal/context/source_system_test.go +++ b/internal/context/source_system_test.go @@ -10,12 +10,7 @@ import ( func TestCollectSystemStateHandlesGitUnavailable(t *testing.T) { t.Parallel() - state, err := collectSystemState(context.Background(), Metadata{ - Workdir: "/workspace", - Shell: "bash", - Provider: "openai", - Model: "gpt-5.4", - }, func(ctx context.Context, workdir string, args ...string) (string, error) { + state, err := collectSystemState(context.Background(), testMetadata("/workspace"), func(ctx context.Context, workdir string, args ...string) (string, error) { return "", errors.New("git unavailable") }) if err != nil { @@ -46,12 +41,7 @@ func TestCollectSystemStateIncludesGitSummary(t *testing.T) { } } - state, err := collectSystemState(context.Background(), Metadata{ - Workdir: "/workspace", - Shell: "bash", - Provider: "openai", - Model: "gpt-5.4", - }, runner) + state, err := collectSystemState(context.Background(), testMetadata("/workspace"), runner) if err != nil { t.Fatalf("collectSystemState() error = %v", err) } @@ -81,9 +71,7 @@ func TestCollectSystemStateReturnsContextError(t *testing.T) { ctx, cancel := context.WithCancel(context.Background()) cancel() - _, err := collectSystemState(ctx, Metadata{ - Workdir: "/workspace", - }, func(ctx context.Context, workdir string, args ...string) (string, error) { + _, err := collectSystemState(ctx, testMetadata("/workspace"), func(ctx context.Context, workdir string, args ...string) (string, error) { return "", ctx.Err() }) if !errors.Is(err, context.Canceled) { diff --git a/internal/context/test_helpers_test.go b/internal/context/test_helpers_test.go new file mode 100644 index 00000000..aeef7728 --- /dev/null +++ b/internal/context/test_helpers_test.go @@ -0,0 +1,13 @@ +package context + +import "neo-code/internal/provider/builtin" + +func testMetadata(workdir string) Metadata { + cfg := builtin.DefaultConfig() + return Metadata{ + Workdir: workdir, + Shell: cfg.Shell, + Provider: cfg.SelectedProvider, + Model: cfg.CurrentModel, + } +}