From 9b58039c9438ceee4e56c5a86e744148d5ae391b Mon Sep 17 00:00:00 2001 From: Yumiue <229866007@qq.com> Date: Thu, 2 Apr 2026 22:30:13 +0800 Subject: [PATCH 01/23] =?UTF-8?q?=EF=BB=BFfeat(compact):=20=E5=AE=8C?= =?UTF-8?q?=E6=88=90=20M1=20=E6=94=B6=E5=8F=A3=E5=B9=B6=E4=BF=AE=E5=A4=8D?= =?UTF-8?q?=E5=85=B3=E9=94=AE=E9=A3=8E=E9=99=A9?= MIME-Version: 1.0 Content-Type: text/plain; charset=UTF-8 Content-Transfer-Encoding: 8bit - 默认启用 micro_compact，新增 manual full_replace 策略 - 统一 compact 流程并补齐 transcript 持久化与唯一命名 - 修复 micro_enabled=false 持久化丢失问题，兼容缺失字段回退默认值 - 调整 TUI compact 轻量可观测，避免 micro 失败污染主流程错误态 - 同步配置/事件/持久化相关文档并补齐回归测试 测试： - go test ./... --- .gitignore | 1 + README.md | 3 + docs/config-management-detail-design.md | 11 + docs/runtime-provider-event-flow.md | 31 +- docs/session-persistence-design.md | 6 + docs/tools-and-tui-integration.md | 2 + internal/config/config_test.go | 178 +++++++++ internal/config/loader.go | 64 ++- internal/config/model.go | 123 +++++- internal/context/compact/runner.go | 500 ++++++++++++++++++++++++ internal/context/compact/runner_test.go | 337 ++++++++++++++++ internal/runtime/events.go | 8 + internal/runtime/runtime.go | 145 ++++++- internal/runtime/runtime_test.go | 231 +++++++++++ internal/tui/commands.go | 4 + internal/tui/commands_test.go | 1 + internal/tui/update.go | 85 +++- internal/tui/update_test.go | 147 ++++++- 18 files changed, 1844 insertions(+), 33 deletions(-) create mode 100644 internal/context/compact/runner.go create mode 100644 internal/context/compact/runner_test.go diff --git a/.gitignore b/.gitignore index df09beca..8c4a7761 100644 --- a/.gitignore +++ b/.gitignore @@ -31,6 +31,7 @@ config.local.yaml # Local data data/ .cache/ +.neocode/projects/**/.transcripts/ # Editor/IDE .idea/ diff --git a/README.md b/README.md index bd65dfe8..17313766 100644 --- a/README.md +++ b/README.md @@ -68,6 +68,9 @@ go run ./cmd/neocode - `/provider` — 切换模型提供商 - `/model` — 切换模型 +- `/compact` — 手动压缩当前会话上下文（保留可追溯 transcript） + +上下文压缩默认开启每轮 `micro_compact`，手动 `/compact` 支持 `keep_recent` 与 `full_replace` 两种策略。 ## 架构概览 diff --git a/docs/config-management-detail-design.md b/docs/config-management-detail-design.md index 8975eb90..9e79bb52 100644 --- a/docs/config-management-detail-design.md +++ b/docs/config-management-detail-design.md @@ -8,6 +8,17 @@ ## 核心类型 - `Config`：顶层应用配置，运行时包含 Provider 列表、当前选中 Provider、当前模型、工作目录、Shell 和循环限制等信息 +- `Config.Context.Compact`：上下文压缩配置，包含： + - `micro_enabled` + - `tool_result_keep_recent` + - `tool_result_placeholder_min_chars` + - `manual_strategy`（M1 支持 `keep_recent` / `full_replace`） + - `manual_keep_recent_spans` + - `max_summary_chars` + +默认策略说明： +- `micro_enabled` 默认 `true`，可通过配置显式关闭 +- `manual_strategy` 默认 `keep_recent` - `ProviderConfig`：单个 Provider 的内建定义，包括 Base URL、默认模型、实例级模型列表和 API Key 环境变量名 - `Manager`：使用 `sync.RWMutex` 保护的配置访问器与修改器 - `Loader`：对 YAML 文件和托管 `.env` 文件的文件系统封装 diff --git a/docs/runtime-provider-event-flow.md b/docs/runtime-provider-event-flow.md index e9e007e2..83a15e37 100644 --- a/docs/runtime-provider-event-flow.md +++ b/docs/runtime-provider-event-flow.md @@ -9,18 +9,23 @@ - `tool_start` - `tool_result` - `error` +- `compact_start` +- `compact_done` +- `compact_error` +- `micro_compact_applied` ## ReAct 主循环 1. 加载目标会话或创建新会话。 2. 追加最新的用户消息。 3. 读取最新配置快照。 -4. 解析当前 provider 配置并构建 provider 实例。 -5. 调用 `context.Builder` 生成本轮请求使用的 `system prompt` 和消息上下文。 -6. 调用 `Provider.Chat`，并把流式事件桥接给 TUI。 -7. 保存 assistant 完整回复。 -8. 执行返回的工具调用，并保存每一个工具结果。 -9. 如果仍需继续推理，则进入下一轮；否则结束。 +4. 每轮请求前先尝试执行 `micro_compact`（失败时降级跳过，不阻断主链路）。 +5. 解析当前 provider 配置并构建 provider 实例。 +6. 调用 `context.Builder` 生成本轮请求使用的 `system prompt` 和消息上下文。 +7. 调用 `Provider.Chat`，并把流式事件桥接给 TUI。 +8. 保存 assistant 完整回复。 +9. 执行返回的工具调用，并保存每一个工具结果。 +10. 如果仍需继续推理，则进入下一轮；否则结束。 ### Context Builder 输入与职责 @@ -59,6 +64,20 @@ - runtime 将其转换成 `RuntimeEvent` - TUI 使用 Bubble Tea `Cmd` 监听事件，并在处理完成后继续订阅 +### Compact 事件 + +- `compact_start`：开始执行 compact（payload 为触发模式，`micro` 或 `manual`） +- `compact_done`：compact 完成，payload 包含： + - `applied` + - `before_chars` + - `after_chars` + - `saved_ratio` + - `trigger_mode` + - `transcript_id` + - `transcript_path` +- `compact_error`：compact 失败，payload 包含触发模式和错误信息 +- `micro_compact_applied`：仅当 `micro_compact` 实际改写了消息时触发 + ## 持久化时机 - 用户消息提交后保存 diff --git a/docs/session-persistence-design.md b/docs/session-persistence-design.md index 0c5d1887..4f5d1cde 100644 --- a/docs/session-persistence-design.md +++ b/docs/session-persistence-design.md @@ -2,14 +2,18 @@ ## 存储策略 NeoCode 在 MVP 阶段使用 JSON 文件持久化 Session，以保持本地优先、易于调试和跨平台可移植。 +除会话主存储外，compact 会在压缩前额外写入 transcript 留痕，确保“移出活跃上下文”而非“不可恢复删除”。 + ## 数据模型 - `Session`：完整消息历史以及 `id`、`title`、`updated_at` 等元信息 - `SessionSummary`：用于侧边栏的轻量摘要结构 +- `Transcript(.jsonl)`：compact 前完整消息快照，每行一条消息，包含 `role/content/tool_calls/tool_call_id/is_error/index/timestamp` ## 加载策略 - `ListSummaries` 只读取渲染侧边栏所需的基础信息 - `Load` 仅在用户真正进入某个会话时读取完整消息历史 - `Save` 通过临时文件原子写入完整 Session +- transcript 当前不参与自动加载回放，仅用于追溯与审计 ## 命名策略 - 新会话默认展示为 `Draft` @@ -18,3 +22,5 @@ NeoCode 在 MVP 阶段使用 JSON 文件持久化 Session，以保持本地优 ## 并发约束 - SessionStore 实现必须自行保护共享访问 - 真正的保存时机由 runtime 决定，TUI 不负责直接触发磁盘写入 +- compact transcript 也通过原子落盘写入，路径为 `~/.neocode/projects//.transcripts/` +- transcript 文件在 Unix 系统默认使用 `0600` 权限，减少敏感内容泄露面 diff --git a/docs/tools-and-tui-integration.md b/docs/tools-and-tui-integration.md index b41a38a3..c092536e 100644 --- a/docs/tools-and-tui-integration.md +++ b/docs/tools-and-tui-integration.md @@ -22,6 +22,8 @@ ## TUI 集成方式 - 本地配置操作统一通过 Slash Command 完成，例如 Base URL、API Key 和模型选择 +- `/compact` 提供手动上下文压缩入口，触发 runtime 的 manual compact 流程 +- TUI 对 compact 仅做轻量可观测展示：显示 `compact_done` / `compact_error` 关键信息，避免事件刷屏 - runtime 事件以内联形式渲染到 transcript 中，而不是单独拆出控制台面板 - 工具开始和结束事件会以轻量提示插入聊天流，使交互更沉浸 diff --git a/internal/config/config_test.go b/internal/config/config_test.go index 97a677b9..70fa4761 100644 --- a/internal/config/config_test.go +++ b/internal/config/config_test.go @@ -812,6 +812,184 @@ func TestConstructorsRejectMissingDependencies(t *testing.T) { }) } +func TestCompactConfigDefaultsAndRoundTrip(t *testing.T) { + tempDir := t.TempDir() + loader := NewLoader(tempDir, testDefaultConfig()) + + cfg, err := loader.Load(context.Background()) + if err != nil { + t.Fatalf("Load() error = %v", err) + } + + compactCfg := cfg.Context.Compact + if !compactCfg.MicroEnabled { + t.Fatalf("expected micro_enabled default true") + } + if compactCfg.ToolResultKeepRecent != DefaultCompactToolResultKeepRecent { + t.Fatalf("expected tool_result_keep_recent=%d, got %d", DefaultCompactToolResultKeepRecent, compactCfg.ToolResultKeepRecent) + } + if compactCfg.ToolResultPlaceholderMinChars != DefaultCompactPlaceholderMinChars { + t.Fatalf("expected tool_result_placeholder_min_chars=%d, got %d", DefaultCompactPlaceholderMinChars, compactCfg.ToolResultPlaceholderMinChars) + } + if compactCfg.ManualStrategy != CompactManualStrategyKeepRecent { + t.Fatalf("expected manual strategy %q, got %q", CompactManualStrategyKeepRecent, compactCfg.ManualStrategy) + } + if compactCfg.ManualKeepRecentSpans != DefaultCompactManualKeepRecentSpans { + t.Fatalf("expected manual_keep_recent_spans=%d, got %d", DefaultCompactManualKeepRecentSpans, compactCfg.ManualKeepRecentSpans) + } + if compactCfg.MaxSummaryChars != DefaultCompactMaxSummaryChars { + t.Fatalf("expected max_summary_chars=%d, got %d", DefaultCompactMaxSummaryChars, compactCfg.MaxSummaryChars) + } + + if err := loader.Save(context.Background(), cfg); err != nil { + t.Fatalf("Save() error = %v", err) + } + reloaded, err := loader.Load(context.Background()) + if err != nil { + t.Fatalf("Reload() error = %v", err) + } + + if reloaded.Context.Compact.ManualStrategy != CompactManualStrategyKeepRecent { + t.Fatalf("expected manual strategy to persist, got %q", reloaded.Context.Compact.ManualStrategy) + } +} + +func TestCompactConfigMicroEnabledFalsePersistsAcrossReload(t *testing.T) { + tempDir := t.TempDir() + loader := NewLoader(tempDir, testDefaultConfig()) + + cfg, err := loader.Load(context.Background()) + if err != nil { + t.Fatalf("Load() error = %v", err) + } + cfg.Context.Compact.MicroEnabled = false + + if err := loader.Save(context.Background(), cfg); err != nil { + t.Fatalf("Save() error = %v", err) + } + reloaded, err := loader.Load(context.Background()) + if err != nil { + t.Fatalf("Reload() error = %v", err) + } + if reloaded.Context.Compact.MicroEnabled { + t.Fatalf("expected micro_enabled to persist as false after reload") + } +} + +func TestCompactConfigMissingMicroEnabledFallsBackToDefaultTrue(t *testing.T) { + tempDir := t.TempDir() + loader := NewLoader(tempDir, testDefaultConfig()) + if err := os.MkdirAll(loader.BaseDir(), 0o755); err != nil { + t.Fatalf("create base dir: %v", err) + } + + legacyConfig := []byte(strings.Join([]string{ + "selected_provider: openai", + "current_model: gpt-4.1", + "workdir: " + filepath.ToSlash(tempDir), + "shell: powershell", + "", + }, "\n")) + if err := os.WriteFile(loader.ConfigPath(), legacyConfig, 0o644); err != nil { + t.Fatalf("write legacy config: %v", err) + } + + cfg, err := loader.Load(context.Background()) + if err != nil { + t.Fatalf("Load() error = %v", err) + } + if !cfg.Context.Compact.MicroEnabled { + t.Fatalf("expected missing micro_enabled to fallback to default true") + } +} + +func TestCompactConfigValidateFailures(t *testing.T) { + tests := []struct { + name string + compact CompactConfig + expectErr string + }{ + { + name: "invalid keep recent", + compact: CompactConfig{ + ToolResultKeepRecent: 0, + ToolResultPlaceholderMinChars: 100, + ManualStrategy: CompactManualStrategyKeepRecent, + ManualKeepRecentSpans: 6, + MaxSummaryChars: 1200, + }, + expectErr: "tool_result_keep_recent", + }, + { + name: "invalid placeholder min chars", + compact: CompactConfig{ + ToolResultKeepRecent: 3, + ToolResultPlaceholderMinChars: 0, + ManualStrategy: CompactManualStrategyKeepRecent, + ManualKeepRecentSpans: 6, + MaxSummaryChars: 1200, + }, + expectErr: "tool_result_placeholder_min_chars", + }, + { + name: "invalid manual strategy", + compact: CompactConfig{ + ToolResultKeepRecent: 3, + ToolResultPlaceholderMinChars: 100, + ManualStrategy: "unsupported", + ManualKeepRecentSpans: 6, + MaxSummaryChars: 1200, + }, + expectErr: "manual_strategy", + }, + { + name: "invalid manual keep spans", + compact: CompactConfig{ + ToolResultKeepRecent: 3, + ToolResultPlaceholderMinChars: 100, + ManualStrategy: CompactManualStrategyKeepRecent, + ManualKeepRecentSpans: 0, + MaxSummaryChars: 1200, + }, + expectErr: "manual_keep_recent_spans", + }, + { + name: "invalid summary chars", + compact: CompactConfig{ + ToolResultKeepRecent: 3, + ToolResultPlaceholderMinChars: 100, + ManualStrategy: CompactManualStrategyKeepRecent, + ManualKeepRecentSpans: 6, + MaxSummaryChars: 0, + }, + expectErr: "max_summary_chars", + }, + } + + for _, tt := range tests { + tt := tt + t.Run(tt.name, func(t *testing.T) { + err := tt.compact.Validate() + if err == nil || !strings.Contains(err.Error(), tt.expectErr) { + t.Fatalf("expected error containing %q, got %v", tt.expectErr, err) + } + }) + } +} + +func TestCompactConfigValidateSupportsFullReplace(t *testing.T) { + err := (CompactConfig{ + ToolResultKeepRecent: 3, + ToolResultPlaceholderMinChars: 100, + ManualStrategy: CompactManualStrategyFullReplace, + ManualKeepRecentSpans: 6, + MaxSummaryChars: 1200, + }).Validate() + if err != nil { + t.Fatalf("expected full_replace strategy to be valid, got %v", err) + } +} + func restoreEnv(t *testing.T, key string) { t.Helper() value, ok := os.LookupEnv(key) diff --git a/internal/config/loader.go b/internal/config/loader.go index 5973f662..365a1508 100644 --- a/internal/config/loader.go +++ b/internal/config/loader.go @@ -22,13 +22,27 @@ type Loader struct { } type persistedConfig struct { - SelectedProvider string `yaml:"selected_provider"` - CurrentModel string `yaml:"current_model"` - Workdir string `yaml:"workdir"` - Shell string `yaml:"shell"` - MaxLoops int `yaml:"max_loops,omitempty"` - ToolTimeoutSec int `yaml:"tool_timeout_sec,omitempty"` - Tools ToolsConfig `yaml:"tools,omitempty"` + SelectedProvider string `yaml:"selected_provider"` + CurrentModel string `yaml:"current_model"` + Workdir string `yaml:"workdir"` + Shell string `yaml:"shell"` + MaxLoops int `yaml:"max_loops,omitempty"` + ToolTimeoutSec int `yaml:"tool_timeout_sec,omitempty"` + Context persistedContextConfig `yaml:"context,omitempty"` + Tools ToolsConfig `yaml:"tools,omitempty"` +} + +type persistedContextConfig struct { + Compact persistedCompactConfig `yaml:"compact,omitempty"` +} + +type persistedCompactConfig struct { + MicroEnabled *bool `yaml:"micro_enabled,omitempty"` + ToolResultKeepRecent int `yaml:"tool_result_keep_recent,omitempty"` + ToolResultPlaceholderMinChars int `yaml:"tool_result_placeholder_min_chars,omitempty"` + ManualStrategy string `yaml:"manual_strategy,omitempty"` + ManualKeepRecentSpans int `yaml:"manual_keep_recent_spans,omitempty"` + MaxSummaryChars int `yaml:"max_summary_chars,omitempty"` } func NewLoader(baseDir string, defaults *Config) *Loader { @@ -122,6 +136,7 @@ func (l *Loader) Save(ctx context.Context, cfg *Config) error { Shell: snapshot.Shell, MaxLoops: snapshot.MaxLoops, ToolTimeoutSec: snapshot.ToolTimeoutSec, + Context: newPersistedContextConfig(snapshot.Context), Tools: snapshot.Tools, } @@ -162,7 +177,7 @@ type aliasConfig struct { WorkspaceRoot string `yaml:"workspace_root"` } -func parseCurrentConfig(data []byte, _ Config) (*Config, error) { +func parseCurrentConfig(data []byte, defaults Config) (*Config, error) { var file persistedConfig if err := yaml.Unmarshal(data, &file); err != nil { return nil, err @@ -185,12 +200,45 @@ func parseCurrentConfig(data []byte, _ Config) (*Config, error) { Shell: strings.TrimSpace(file.Shell), MaxLoops: file.MaxLoops, ToolTimeoutSec: file.ToolTimeoutSec, + Context: fromPersistedContextConfig(file.Context, defaults.Context), Tools: file.Tools, } return cfg, nil } +func newPersistedContextConfig(cfg ContextConfig) persistedContextConfig { + microEnabled := cfg.Compact.MicroEnabled + return persistedContextConfig{ + Compact: persistedCompactConfig{ + MicroEnabled: µEnabled, + ToolResultKeepRecent: cfg.Compact.ToolResultKeepRecent, + ToolResultPlaceholderMinChars: cfg.Compact.ToolResultPlaceholderMinChars, + ManualStrategy: cfg.Compact.ManualStrategy, + ManualKeepRecentSpans: cfg.Compact.ManualKeepRecentSpans, + MaxSummaryChars: cfg.Compact.MaxSummaryChars, + }, + } +} + +func fromPersistedContextConfig(file persistedContextConfig, defaults ContextConfig) ContextConfig { + out := ContextConfig{ + Compact: CompactConfig{ + ToolResultKeepRecent: file.Compact.ToolResultKeepRecent, + ToolResultPlaceholderMinChars: file.Compact.ToolResultPlaceholderMinChars, + ManualStrategy: strings.TrimSpace(file.Compact.ManualStrategy), + ManualKeepRecentSpans: file.Compact.ManualKeepRecentSpans, + MaxSummaryChars: file.Compact.MaxSummaryChars, + }, + } + if file.Compact.MicroEnabled != nil { + out.Compact.MicroEnabled = *file.Compact.MicroEnabled + } else { + out.Compact.MicroEnabled = defaults.Compact.MicroEnabled + } + return out +} + func requiresConfigRewrite(data []byte) bool { text := strings.TrimSpace(string(data)) switch { diff --git a/internal/config/model.go b/internal/config/model.go index c0b57f0b..274ba096 100644 --- a/internal/config/model.go +++ b/internal/config/model.go @@ -11,10 +11,19 @@ import ( ) const ( - DefaultWorkdir = "." - DefaultMaxLoops = 8 - DefaultToolTimeoutSec = 20 - DefaultWebFetchMaxResponseBytes int64 = 256 * 1024 + DefaultWorkdir = "." + DefaultMaxLoops = 8 + DefaultToolTimeoutSec = 20 + DefaultWebFetchMaxResponseBytes int64 = 256 * 1024 + DefaultCompactToolResultKeepRecent = 3 + DefaultCompactPlaceholderMinChars = 100 + DefaultCompactManualKeepRecentSpans = 6 + DefaultCompactMaxSummaryChars = 1200 +) + +const ( + CompactManualStrategyKeepRecent = "keep_recent" + CompactManualStrategyFullReplace = "full_replace" ) var defaultWebFetchSupportedContentTypes = []string{ @@ -34,6 +43,7 @@ type Config struct { Shell string `yaml:"shell"` MaxLoops int `yaml:"max_loops,omitempty"` ToolTimeoutSec int `yaml:"tool_timeout_sec,omitempty"` + Context ContextConfig `yaml:"context,omitempty"` Tools ToolsConfig `yaml:"tools,omitempty"` } @@ -55,6 +65,19 @@ type ToolsConfig struct { WebFetch WebFetchConfig `yaml:"webfetch,omitempty"` } +type ContextConfig struct { + Compact CompactConfig `yaml:"compact,omitempty"` +} + +type CompactConfig struct { + MicroEnabled bool `yaml:"micro_enabled,omitempty"` + ToolResultKeepRecent int `yaml:"tool_result_keep_recent,omitempty"` + ToolResultPlaceholderMinChars int `yaml:"tool_result_placeholder_min_chars,omitempty"` + ManualStrategy string `yaml:"manual_strategy,omitempty"` + ManualKeepRecentSpans int `yaml:"manual_keep_recent_spans,omitempty"` + MaxSummaryChars int `yaml:"max_summary_chars,omitempty"` +} + type WebFetchConfig struct { MaxResponseBytes int64 `yaml:"max_response_bytes,omitempty"` SupportedContentTypes []string `yaml:"supported_content_types,omitempty"` @@ -70,6 +93,7 @@ func Default() *Config { Shell: defaultShell(), MaxLoops: DefaultMaxLoops, ToolTimeoutSec: DefaultToolTimeoutSec, + Context: defaultContextConfig(), Tools: ToolsConfig{ WebFetch: defaultWebFetchConfig(), }, @@ -83,6 +107,7 @@ func (c *Config) Clone() Config { clone := *c clone.Providers = cloneProviders(c.Providers) + clone.Context = c.Context.Clone() clone.Tools = c.Tools.Clone() return clone } @@ -123,6 +148,7 @@ func (c *Config) ApplyDefaultsFrom(defaults Config) { if c.ToolTimeoutSec <= 0 { c.ToolTimeoutSec = defaults.ToolTimeoutSec } + c.Context.ApplyDefaults(defaults.Context) c.Tools.ApplyDefaults(defaults.Tools) c.Workdir = normalizeWorkdir(c.Workdir) @@ -174,6 +200,9 @@ func (c *Config) Validate() error { if err := c.Tools.Validate(); err != nil { return fmt.Errorf("config: tools: %w", err) } + if err := c.Context.Validate(); err != nil { + return fmt.Errorf("config: context: %w", err) + } return nil } @@ -355,12 +384,35 @@ func defaultWebFetchConfig() WebFetchConfig { } } +func defaultContextConfig() ContextConfig { + return ContextConfig{ + Compact: defaultCompactConfig(), + } +} + +func defaultCompactConfig() CompactConfig { + return CompactConfig{ + MicroEnabled: true, + ToolResultKeepRecent: DefaultCompactToolResultKeepRecent, + ToolResultPlaceholderMinChars: DefaultCompactPlaceholderMinChars, + ManualStrategy: CompactManualStrategyKeepRecent, + ManualKeepRecentSpans: DefaultCompactManualKeepRecentSpans, + MaxSummaryChars: DefaultCompactMaxSummaryChars, + } +} + func (c ToolsConfig) Clone() ToolsConfig { return ToolsConfig{ WebFetch: c.WebFetch.Clone(), } } +func (c ContextConfig) Clone() ContextConfig { + return ContextConfig{ + Compact: c.Compact.Clone(), + } +} + func (c *ToolsConfig) ApplyDefaults(defaults ToolsConfig) { if c == nil { return @@ -369,6 +421,14 @@ func (c *ToolsConfig) ApplyDefaults(defaults ToolsConfig) { c.WebFetch.ApplyDefaults(defaults.WebFetch) } +func (c *ContextConfig) ApplyDefaults(defaults ContextConfig) { + if c == nil { + return + } + + c.Compact.ApplyDefaults(defaults.Compact) +} + func (c ToolsConfig) Validate() error { if err := c.WebFetch.Validate(); err != nil { return fmt.Errorf("webfetch: %w", err) @@ -376,12 +436,23 @@ func (c ToolsConfig) Validate() error { return nil } +func (c ContextConfig) Validate() error { + if err := c.Compact.Validate(); err != nil { + return fmt.Errorf("compact: %w", err) + } + return nil +} + func (c WebFetchConfig) Clone() WebFetchConfig { clone := c clone.SupportedContentTypes = append([]string(nil), c.SupportedContentTypes...) return clone } +func (c CompactConfig) Clone() CompactConfig { + return c +} + func (c *WebFetchConfig) ApplyDefaults(defaults WebFetchConfig) { if c == nil { return @@ -393,6 +464,28 @@ func (c *WebFetchConfig) ApplyDefaults(defaults WebFetchConfig) { c.SupportedContentTypes = normalizeContentTypes(c.SupportedContentTypes, defaults.SupportedContentTypes) } +func (c *CompactConfig) ApplyDefaults(defaults CompactConfig) { + if c == nil { + return + } + + if c.ToolResultKeepRecent <= 0 { + c.ToolResultKeepRecent = defaults.ToolResultKeepRecent + } + if c.ToolResultPlaceholderMinChars <= 0 { + c.ToolResultPlaceholderMinChars = defaults.ToolResultPlaceholderMinChars + } + if strings.TrimSpace(c.ManualStrategy) == "" { + c.ManualStrategy = defaults.ManualStrategy + } + if c.ManualKeepRecentSpans <= 0 { + c.ManualKeepRecentSpans = defaults.ManualKeepRecentSpans + } + if c.MaxSummaryChars <= 0 { + c.MaxSummaryChars = defaults.MaxSummaryChars + } +} + func (c WebFetchConfig) Validate() error { if c.MaxResponseBytes <= 0 { return errors.New("max_response_bytes must be greater than 0") @@ -409,6 +502,28 @@ func (c WebFetchConfig) Validate() error { return nil } +func (c CompactConfig) Validate() error { + if c.ToolResultKeepRecent <= 0 { + return errors.New("tool_result_keep_recent must be greater than 0") + } + if c.ToolResultPlaceholderMinChars <= 0 { + return errors.New("tool_result_placeholder_min_chars must be greater than 0") + } + if c.ManualKeepRecentSpans <= 0 { + return errors.New("manual_keep_recent_spans must be greater than 0") + } + if c.MaxSummaryChars <= 0 { + return errors.New("max_summary_chars must be greater than 0") + } + + switch strings.ToLower(strings.TrimSpace(c.ManualStrategy)) { + case CompactManualStrategyKeepRecent, CompactManualStrategyFullReplace: + return nil + default: + return fmt.Errorf("manual_strategy %q is not supported", c.ManualStrategy) + } +} + func normalizeContentTypes(values []string, defaults []string) []string { source := values if len(source) == 0 { diff --git a/internal/context/compact/runner.go b/internal/context/compact/runner.go new file mode 100644 index 00000000..2702892d --- /dev/null +++ b/internal/context/compact/runner.go @@ -0,0 +1,500 @@ +package compact + +import ( + "context" + cryptorand "crypto/rand" + "crypto/sha1" + "encoding/hex" + "encoding/json" + "errors" + "fmt" + "os" + "path/filepath" + "regexp" + goruntime "runtime" + "strings" + "time" + + "neo-code/internal/config" + "neo-code/internal/provider" +) + +type Mode string + +const ( + ModeMicro Mode = "micro" + ModeManual Mode = "manual" +) + +type ErrorMode string + +const ( + ErrorModeNone ErrorMode = "none" +) + +type Input struct { + Mode Mode + SessionID string + Workdir string + Messages []provider.Message + Config config.CompactConfig +} + +type Metrics struct { + BeforeChars int `json:"before_chars"` + AfterChars int `json:"after_chars"` + SavedRatio float64 `json:"saved_ratio"` + TriggerMode string `json:"trigger_mode"` +} + +type Result struct { + Messages []provider.Message `json:"messages"` + Metrics Metrics `json:"metrics"` + TranscriptID string `json:"transcript_id"` + TranscriptPath string `json:"transcript_path"` + Applied bool `json:"applied"` + ErrorMode ErrorMode `json:"error_mode"` +} + +type Runner interface { + Run(ctx context.Context, input Input) (Result, error) +} + +type Service struct { + now func() time.Time + randomToken func() (string, error) + userHomeDir func() (string, error) + mkdirAll func(path string, perm os.FileMode) error + writeFile func(name string, data []byte, perm os.FileMode) error + rename func(oldPath, newPath string) error + remove func(path string) error +} + +func NewRunner() *Service { + return &Service{ + now: time.Now, + randomToken: randomTranscriptToken, + userHomeDir: os.UserHomeDir, + mkdirAll: os.MkdirAll, + writeFile: os.WriteFile, + rename: os.Rename, + remove: os.Remove, + } +} + +func (s *Service) Run(ctx context.Context, input Input) (Result, error) { + if err := ctx.Err(); err != nil { + return Result{}, err + } + + cfg := normalizeCompactConfig(input.Config) + messages := cloneMessages(input.Messages) + + beforeChars := countMessageChars(messages) + base := Result{ + Messages: messages, + Applied: false, + ErrorMode: ErrorModeNone, + Metrics: Metrics{ + BeforeChars: beforeChars, + AfterChars: beforeChars, + SavedRatio: 0, + TriggerMode: string(input.Mode), + }, + } + + switch input.Mode { + case ModeMicro: + if !cfg.MicroEnabled || !hasMicroCandidate(messages, cfg) { + return base, nil + } + case ModeManual: + // manual compact is always evaluated when explicitly requested. + default: + return Result{}, fmt.Errorf("compact: unsupported mode %q", input.Mode) + } + + transcriptID, transcriptPath, err := s.saveTranscript(messages, strings.TrimSpace(input.SessionID), strings.TrimSpace(input.Workdir)) + if err != nil { + return Result{}, err + } + base.TranscriptID = transcriptID + base.TranscriptPath = transcriptPath + + var ( + next []provider.Message + applied bool + ) + + switch input.Mode { + case ModeMicro: + next, applied = microCompact(messages, cfg) + case ModeManual: + next, applied, err = manualCompact(messages, cfg) + if err != nil { + return Result{}, err + } + } + + afterChars := countMessageChars(next) + result := base + result.Messages = next + result.Applied = applied + result.Metrics.AfterChars = afterChars + if beforeChars > 0 { + result.Metrics.SavedRatio = float64(beforeChars-afterChars) / float64(beforeChars) + } + return result, nil +} + +func hasMicroCandidate(messages []provider.Message, cfg config.CompactConfig) bool { + toolIndices := make([]int, 0, len(messages)) + for i, message := range messages { + if message.Role == provider.RoleTool { + toolIndices = append(toolIndices, i) + } + } + if len(toolIndices) <= cfg.ToolResultKeepRecent { + return false + } + candidateCount := len(toolIndices) - cfg.ToolResultKeepRecent + for i := 0; i < candidateCount; i++ { + if len(messages[toolIndices[i]].Content) >= cfg.ToolResultPlaceholderMinChars { + return true + } + } + return false +} + +func microCompact(messages []provider.Message, cfg config.CompactConfig) ([]provider.Message, bool) { + next := cloneMessages(messages) + + toolIndices := make([]int, 0, len(next)) + for i, message := range next { + if message.Role == provider.RoleTool { + toolIndices = append(toolIndices, i) + } + } + + if len(toolIndices) <= cfg.ToolResultKeepRecent { + return next, false + } + + toolNameByCallID := buildToolNameIndex(next) + candidateCount := len(toolIndices) - cfg.ToolResultKeepRecent + applied := false + for i := 0; i < candidateCount; i++ { + idx := toolIndices[i] + message := next[idx] + if len(message.Content) < cfg.ToolResultPlaceholderMinChars { + continue + } + + toolName := strings.TrimSpace(toolNameByCallID[message.ToolCallID]) + if toolName == "" { + toolName = "unknown_tool" + } + placeholder := fmt.Sprintf("[Previous tool used: %s]", toolName) + if message.Content == placeholder { + continue + } + message.Content = placeholder + next[idx] = message + applied = true + } + + return next, applied +} + +func buildToolNameIndex(messages []provider.Message) map[string]string { + index := make(map[string]string) + for _, message := range messages { + if message.Role != provider.RoleAssistant || len(message.ToolCalls) == 0 { + continue + } + for _, call := range message.ToolCalls { + id := strings.TrimSpace(call.ID) + name := strings.TrimSpace(call.Name) + if id == "" || name == "" { + continue + } + index[id] = name + } + } + return index +} + +type span struct { + start int + end int +} + +func collectSpans(messages []provider.Message) []span { + spans := make([]span, 0, len(messages)) + for i := 0; i < len(messages); { + start := i + i++ + + if messages[start].Role == provider.RoleAssistant && len(messages[start].ToolCalls) > 0 { + for i < len(messages) && messages[i].Role == provider.RoleTool { + i++ + } + } + + spans = append(spans, span{start: start, end: i}) + } + return spans +} + +func manualCompact(messages []provider.Message, cfg config.CompactConfig) ([]provider.Message, bool, error) { + strategy := strings.ToLower(strings.TrimSpace(cfg.ManualStrategy)) + switch strategy { + case config.CompactManualStrategyKeepRecent: + return manualCompactKeepRecent(messages, cfg) + case config.CompactManualStrategyFullReplace: + return manualCompactFullReplace(messages, cfg) + default: + return nil, false, fmt.Errorf("compact: manual strategy %q is not supported", cfg.ManualStrategy) + } +} + +func manualCompactKeepRecent(messages []provider.Message, cfg config.CompactConfig) ([]provider.Message, bool, error) { + spans := collectSpans(messages) + if len(spans) <= cfg.ManualKeepRecentSpans { + return cloneMessages(messages), false, nil + } + + keepStart := spans[len(spans)-cfg.ManualKeepRecentSpans].start + removed := cloneMessages(messages[:keepStart]) + kept := cloneMessages(messages[keepStart:]) + + summary, err := buildSummary(removed, len(spans)-cfg.ManualKeepRecentSpans, cfg) + if err != nil { + return nil, false, err + } + + next := make([]provider.Message, 0, len(kept)+1) + next = append(next, provider.Message{Role: provider.RoleAssistant, Content: summary}) + next = append(next, kept...) + return next, true, nil +} + +func manualCompactFullReplace(messages []provider.Message, cfg config.CompactConfig) ([]provider.Message, bool, error) { + if len(messages) == 0 { + return nil, false, nil + } + spans := collectSpans(messages) + summary, err := buildSummary(cloneMessages(messages), len(spans), cfg) + if err != nil { + return nil, false, err + } + + return []provider.Message{{Role: provider.RoleAssistant, Content: summary}}, true, nil +} + +func buildSummary(removed []provider.Message, removedSpans int, cfg config.CompactConfig) (string, error) { + toolNames := make([]string, 0, 8) + seenTools := map[string]struct{}{} + for _, message := range removed { + for _, call := range message.ToolCalls { + name := strings.TrimSpace(call.Name) + if name == "" { + continue + } + if _, exists := seenTools[name]; exists { + continue + } + seenTools[name] = struct{}{} + toolNames = append(toolNames, name) + } + } + toolSummary := "none" + if len(toolNames) > 0 { + toolSummary = strings.Join(toolNames, ", ") + } + + summary := fmt.Sprintf(strings.Join([]string{ + "[compact_summary]", + "done:", + "- Archived %d historical spans (%d messages).", + "", + "in_progress:", + "- Continue from the retained recent context window.", + "", + "decisions:", + "- manual_strategy=%s", + "- manual_keep_recent_spans=%d", + "", + "code_changes:", + "- Older context outside the recent window was replaced by this summary.", + "- Historical tool calls in archived spans: %s", + "", + "constraints:", + "- Assistant tool_calls and tool_result pairs remain intact in retained spans.", + }, "\n"), removedSpans, len(removed), cfg.ManualStrategy, cfg.ManualKeepRecentSpans, toolSummary) + + return validateSummary(summary, cfg.MaxSummaryChars) +} + +func validateSummary(summary string, maxChars int) (string, error) { + summary = strings.TrimSpace(summary) + if maxChars > 0 && len(summary) > maxChars { + summary = strings.TrimSpace(summary[:maxChars]) + } + + hasDone := sectionHasContent(summary, "done") + hasInProgress := sectionHasContent(summary, "in_progress") + if !hasDone && !hasInProgress { + return "", errors.New("compact: summary requires done or in_progress content") + } + return summary, nil +} + +func sectionHasContent(summary string, section string) bool { + pattern := fmt.Sprintf(`(?ms)^%s:\s*\n\s*-\s+.+?(\n\w+?:|\z)`, regexp.QuoteMeta(section)) + re, err := regexp.Compile(pattern) + if err != nil { + return false + } + return re.MatchString(summary) +} + +type transcriptLine struct { + Index int `json:"index"` + Timestamp string `json:"timestamp"` + Role string `json:"role"` + Content string `json:"content"` + ToolCalls []provider.ToolCall `json:"tool_calls,omitempty"` + ToolCallID string `json:"tool_call_id,omitempty"` + IsError bool `json:"is_error,omitempty"` +} + +func (s *Service) saveTranscript(messages []provider.Message, sessionID string, workdir string) (string, string, error) { + home, err := s.userHomeDir() + if err != nil { + return "", "", fmt.Errorf("compact: resolve user home: %w", err) + } + + projectHash := hashProject(workdir) + dir := filepath.Join(home, ".neocode", "projects", projectHash, ".transcripts") + if err := s.mkdirAll(dir, 0o755); err != nil { + return "", "", fmt.Errorf("compact: create transcript dir: %w", err) + } + + sessionID = sanitizeID(sessionID) + if sessionID == "" { + sessionID = "draft" + } + tokenFn := s.randomToken + if tokenFn == nil { + tokenFn = randomTranscriptToken + } + randomToken, err := tokenFn() + if err != nil { + return "", "", fmt.Errorf("compact: generate transcript token: %w", err) + } + transcriptID := fmt.Sprintf("transcript_%d_%s_%s", s.now().UnixNano(), randomToken, sessionID) + transcriptPath := filepath.Join(dir, transcriptID+".jsonl") + tmpPath := transcriptPath + ".tmp" + + now := s.now().UTC().Format(time.RFC3339Nano) + var builder strings.Builder + for i, message := range messages { + line := transcriptLine{ + Index: i, + Timestamp: now, + Role: message.Role, + Content: message.Content, + ToolCalls: append([]provider.ToolCall(nil), message.ToolCalls...), + ToolCallID: message.ToolCallID, + IsError: message.IsError, + } + payload, err := json.Marshal(line) + if err != nil { + return "", "", fmt.Errorf("compact: marshal transcript line: %w", err) + } + builder.Write(payload) + builder.WriteByte('\n') + } + + if err := s.writeFile(tmpPath, []byte(builder.String()), transcriptFileMode()); err != nil { + return "", "", fmt.Errorf("compact: write transcript: %w", err) + } + if err := s.rename(tmpPath, transcriptPath); err != nil { + _ = s.remove(tmpPath) + return "", "", fmt.Errorf("compact: commit transcript: %w", err) + } + + return transcriptID, transcriptPath, nil +} + +func transcriptFileMode() os.FileMode { + if goruntime.GOOS == "windows" { + return 0o644 + } + return 0o600 +} + +func randomTranscriptToken() (string, error) { + entropy := make([]byte, 4) + if _, err := cryptorand.Read(entropy); err != nil { + return "", err + } + return hex.EncodeToString(entropy), nil +} + +func hashProject(workdir string) string { + clean := strings.TrimSpace(filepath.Clean(workdir)) + if clean == "" { + clean = "unknown" + } + sum := sha1.Sum([]byte(strings.ToLower(clean))) + return hex.EncodeToString(sum[:8]) +} + +var nonIDChars = regexp.MustCompile(`[^a-zA-Z0-9_-]+`) + +func sanitizeID(value string) string { + value = strings.TrimSpace(value) + if value == "" { + return "" + } + return nonIDChars.ReplaceAllString(value, "_") +} + +func cloneMessages(messages []provider.Message) []provider.Message { + if len(messages) == 0 { + return nil + } + out := make([]provider.Message, 0, len(messages)) + for _, message := range messages { + next := message + next.ToolCalls = append([]provider.ToolCall(nil), message.ToolCalls...) + out = append(out, next) + } + return out +} + +func countMessageChars(messages []provider.Message) int { + total := 0 + for _, message := range messages { + total += len(message.Role) + total += len(message.Content) + total += len(message.ToolCallID) + for _, call := range message.ToolCalls { + total += len(call.ID) + total += len(call.Name) + total += len(call.Arguments) + } + } + return total +} + +func normalizeCompactConfig(cfg config.CompactConfig) config.CompactConfig { + defaults := config.Default().Context.Compact + cfg.ApplyDefaults(defaults) + if strings.TrimSpace(cfg.ManualStrategy) == "" { + cfg.ManualStrategy = config.CompactManualStrategyKeepRecent + } + return cfg +} diff --git a/internal/context/compact/runner_test.go b/internal/context/compact/runner_test.go new file mode 100644 index 00000000..8b1608f5 --- /dev/null +++ b/internal/context/compact/runner_test.go @@ -0,0 +1,337 @@ +package compact + +import ( + "context" + "errors" + "os" + "path/filepath" + goruntime "runtime" + "strings" + "testing" + "time" + + "neo-code/internal/config" + "neo-code/internal/provider" +) + +func TestMicroCompactReplacesOnlyOldLongToolResults(t *testing.T) { + t.Parallel() + + runner := NewRunner() + home := t.TempDir() + runner.userHomeDir = func() (string, error) { return home, nil } + + messages := []provider.Message{ + {Role: provider.RoleUser, Content: "start"}, + {Role: provider.RoleAssistant, ToolCalls: []provider.ToolCall{{ID: "call-1", Name: "filesystem_read_file", Arguments: "{}"}}}, + {Role: provider.RoleTool, ToolCallID: "call-1", Content: strings.Repeat("A", 40)}, + {Role: provider.RoleAssistant, ToolCalls: []provider.ToolCall{{ID: "call-2", Name: "filesystem_grep", Arguments: "{}"}}}, + {Role: provider.RoleTool, ToolCallID: "call-2", Content: "short"}, + {Role: provider.RoleAssistant, ToolCalls: []provider.ToolCall{{ID: "call-3", Name: "bash", Arguments: "{}"}}}, + {Role: provider.RoleTool, ToolCallID: "call-3", Content: strings.Repeat("B", 50)}, + } + + result, err := runner.Run(context.Background(), Input{ + Mode: ModeMicro, + SessionID: "session-a", + Workdir: t.TempDir(), + Messages: messages, + Config: config.CompactConfig{ + MicroEnabled: true, + ToolResultKeepRecent: 1, + ToolResultPlaceholderMinChars: 10, + ManualStrategy: config.CompactManualStrategyKeepRecent, + ManualKeepRecentSpans: 6, + MaxSummaryChars: 1200, + }, + }) + if err != nil { + t.Fatalf("Run() error = %v", err) + } + if !result.Applied { + t.Fatalf("expected micro compact applied") + } + if !strings.Contains(result.Messages[2].Content, "filesystem_read_file") { + t.Fatalf("expected first old tool result to be replaced, got %q", result.Messages[2].Content) + } + if result.Messages[4].Content != "short" { + t.Fatalf("expected short old tool result unchanged, got %q", result.Messages[4].Content) + } + if result.Messages[6].Content == "[Previous tool used: bash]" { + t.Fatalf("expected recent tool result to be retained") + } + if result.TranscriptID == "" || result.TranscriptPath == "" { + t.Fatalf("expected transcript metadata, got %+v", result) + } + if _, err := os.Stat(result.TranscriptPath); err != nil { + t.Fatalf("expected transcript file: %v", err) + } +} + +func TestMicroCompactFallsBackToUnknownTool(t *testing.T) { + t.Parallel() + + runner := NewRunner() + home := t.TempDir() + runner.userHomeDir = func() (string, error) { return home, nil } + + messages := []provider.Message{ + {Role: provider.RoleUser, Content: "start"}, + {Role: provider.RoleTool, ToolCallID: "missing-call", Content: strings.Repeat("X", 24)}, + {Role: provider.RoleTool, ToolCallID: "recent-call", Content: strings.Repeat("Y", 24)}, + } + + result, err := runner.Run(context.Background(), Input{ + Mode: ModeMicro, + SessionID: "session-b", + Workdir: t.TempDir(), + Messages: messages, + Config: config.CompactConfig{ + MicroEnabled: true, + ToolResultKeepRecent: 1, + ToolResultPlaceholderMinChars: 10, + ManualStrategy: config.CompactManualStrategyKeepRecent, + ManualKeepRecentSpans: 6, + MaxSummaryChars: 1200, + }, + }) + if err != nil { + t.Fatalf("Run() error = %v", err) + } + if got := result.Messages[1].Content; got != "[Previous tool used: unknown_tool]" { + t.Fatalf("expected unknown tool placeholder, got %q", got) + } +} + +func TestManualCompactAddsSummaryAndKeepsRecentSpans(t *testing.T) { + t.Parallel() + + runner := NewRunner() + home := t.TempDir() + runner.userHomeDir = func() (string, error) { return home, nil } + + messages := []provider.Message{ + {Role: provider.RoleUser, Content: "old requirement"}, + {Role: provider.RoleAssistant, ToolCalls: []provider.ToolCall{{ID: "call-old", Name: "filesystem_grep", Arguments: "{}"}}}, + {Role: provider.RoleTool, ToolCallID: "call-old", Content: "old result"}, + {Role: provider.RoleAssistant, Content: "latest answer"}, + } + + result, err := runner.Run(context.Background(), Input{ + Mode: ModeManual, + SessionID: "session-c", + Workdir: t.TempDir(), + Messages: messages, + Config: config.CompactConfig{ + MicroEnabled: false, + ToolResultKeepRecent: 3, + ToolResultPlaceholderMinChars: 100, + ManualStrategy: config.CompactManualStrategyKeepRecent, + ManualKeepRecentSpans: 1, + MaxSummaryChars: 1200, + }, + }) + if err != nil { + t.Fatalf("Run() error = %v", err) + } + if !result.Applied { + t.Fatalf("expected manual compact applied") + } + if len(result.Messages) != 2 { + t.Fatalf("expected summary + 1 kept span, got %d", len(result.Messages)) + } + summary := result.Messages[0] + if summary.Role != provider.RoleAssistant { + t.Fatalf("expected summary role assistant, got %q", summary.Role) + } + for _, section := range []string{"done:", "in_progress:", "decisions:", "code_changes:", "constraints:"} { + if !strings.Contains(summary.Content, section) { + t.Fatalf("expected summary to include section %q, got %q", section, summary.Content) + } + } + if result.Messages[1].Content != "latest answer" { + t.Fatalf("expected newest span kept, got %+v", result.Messages[1]) + } +} + +func TestManualCompactWritesTranscriptJSONL(t *testing.T) { + t.Parallel() + + runner := NewRunner() + home := t.TempDir() + runner.userHomeDir = func() (string, error) { return home, nil } + + result, err := runner.Run(context.Background(), Input{ + Mode: ModeManual, + SessionID: "session-jsonl", + Workdir: filepath.Join(home, "workspace"), + Messages: []provider.Message{ + {Role: provider.RoleUser, Content: "hello"}, + }, + Config: config.CompactConfig{ + ManualStrategy: config.CompactManualStrategyKeepRecent, + ManualKeepRecentSpans: 6, + ToolResultKeepRecent: 3, + ToolResultPlaceholderMinChars: 100, + MaxSummaryChars: 1200, + }, + }) + if err != nil { + t.Fatalf("Run() error = %v", err) + } + data, err := os.ReadFile(result.TranscriptPath) + if err != nil { + t.Fatalf("read transcript: %v", err) + } + if !strings.Contains(string(data), `"role":"user"`) { + t.Fatalf("expected jsonl content, got %q", string(data)) + } + if !strings.Contains(filepath.ToSlash(result.TranscriptPath), "/.neocode/projects/") { + t.Fatalf("expected transcript path under .neocode/projects, got %q", result.TranscriptPath) + } + if !strings.HasPrefix(result.TranscriptID, "transcript_") { + t.Fatalf("unexpected transcript id: %q", result.TranscriptID) + } + if goruntime.GOOS != "windows" { + info, err := os.Stat(result.TranscriptPath) + if err != nil { + t.Fatalf("stat transcript: %v", err) + } + if got := info.Mode().Perm(); got != 0o600 { + t.Fatalf("expected transcript mode 0600, got %04o", got) + } + } +} + +func TestManualCompactFailsWhenTranscriptWriteFails(t *testing.T) { + t.Parallel() + + runner := NewRunner() + runner.userHomeDir = func() (string, error) { return t.TempDir(), nil } + runner.mkdirAll = func(path string, perm os.FileMode) error { + return errors.New("disk full") + } + + _, err := runner.Run(context.Background(), Input{ + Mode: ModeManual, + SessionID: "session-fail", + Workdir: t.TempDir(), + Messages: []provider.Message{ + {Role: provider.RoleUser, Content: "hello"}, + }, + Config: config.CompactConfig{ + ManualStrategy: config.CompactManualStrategyKeepRecent, + ManualKeepRecentSpans: 6, + ToolResultKeepRecent: 3, + ToolResultPlaceholderMinChars: 100, + MaxSummaryChars: 1200, + }, + }) + if err == nil || !strings.Contains(err.Error(), "disk full") { + t.Fatalf("expected transcript write failure, got %v", err) + } +} + +func TestManualCompactFullReplaceRewritesAllMessages(t *testing.T) { + t.Parallel() + + runner := NewRunner() + home := t.TempDir() + runner.userHomeDir = func() (string, error) { return home, nil } + + messages := []provider.Message{ + {Role: provider.RoleUser, Content: "old requirement"}, + {Role: provider.RoleAssistant, ToolCalls: []provider.ToolCall{{ID: "call-old", Name: "filesystem_grep", Arguments: "{}"}}}, + {Role: provider.RoleTool, ToolCallID: "call-old", Content: "old result"}, + {Role: provider.RoleAssistant, Content: "latest answer"}, + } + + result, err := runner.Run(context.Background(), Input{ + Mode: ModeManual, + SessionID: "session-full-replace", + Workdir: t.TempDir(), + Messages: messages, + Config: config.CompactConfig{ + MicroEnabled: true, + ToolResultKeepRecent: 3, + ToolResultPlaceholderMinChars: 100, + ManualStrategy: config.CompactManualStrategyFullReplace, + ManualKeepRecentSpans: 6, + MaxSummaryChars: 1200, + }, + }) + if err != nil { + t.Fatalf("Run() error = %v", err) + } + if !result.Applied { + t.Fatalf("expected full_replace compact applied") + } + if len(result.Messages) != 1 { + t.Fatalf("expected single summary message, got %d", len(result.Messages)) + } + if result.Messages[0].Role != provider.RoleAssistant { + t.Fatalf("expected summary role assistant, got %q", result.Messages[0].Role) + } + for _, section := range []string{"done:", "in_progress:", "decisions:", "code_changes:", "constraints:"} { + if !strings.Contains(result.Messages[0].Content, section) { + t.Fatalf("expected summary section %q, got %q", section, result.Messages[0].Content) + } + } +} + +func TestSaveTranscriptUsesUniqueIDWithinSameTimestamp(t *testing.T) { + t.Parallel() + + runner := NewRunner() + home := t.TempDir() + runner.userHomeDir = func() (string, error) { return home, nil } + fixedNow := time.Unix(1712052000, 123456789) + runner.now = func() time.Time { return fixedNow } + tokenSeq := []string{"a1b2c3d4", "b2c3d4e5"} + runner.randomToken = func() (string, error) { + if len(tokenSeq) == 0 { + return "", errors.New("empty token sequence") + } + next := tokenSeq[0] + tokenSeq = tokenSeq[1:] + return next, nil + } + + input := Input{ + Mode: ModeManual, + SessionID: "session-dup-safe", + Workdir: t.TempDir(), + Messages: []provider.Message{ + {Role: provider.RoleUser, Content: "hello"}, + {Role: provider.RoleAssistant, Content: "world"}, + }, + Config: config.CompactConfig{ + ManualStrategy: config.CompactManualStrategyFullReplace, + ManualKeepRecentSpans: 6, + ToolResultKeepRecent: 3, + ToolResultPlaceholderMinChars: 100, + MaxSummaryChars: 1200, + }, + } + + first, err := runner.Run(context.Background(), input) + if err != nil { + t.Fatalf("first Run() error = %v", err) + } + second, err := runner.Run(context.Background(), input) + if err != nil { + t.Fatalf("second Run() error = %v", err) + } + if first.TranscriptID == second.TranscriptID { + t.Fatalf("expected distinct transcript ids, got %q", first.TranscriptID) + } + if first.TranscriptPath == second.TranscriptPath { + t.Fatalf("expected distinct transcript paths, got %q", first.TranscriptPath) + } + if _, err := os.Stat(first.TranscriptPath); err != nil { + t.Fatalf("first transcript file missing: %v", err) + } + if _, err := os.Stat(second.TranscriptPath); err != nil { + t.Fatalf("second transcript file missing: %v", err) + } +} diff --git a/internal/runtime/events.go b/internal/runtime/events.go index 0e2ce006..27cd73b7 100644 --- a/internal/runtime/events.go +++ b/internal/runtime/events.go @@ -36,4 +36,12 @@ const ( // EventProviderRetry is emitted when runtime retries a provider call due to // a retryable error (e.g. 429, 5xx). Payload is a human-readable message. EventProviderRetry EventType = "provider_retry" + // EventCompactStart is emitted when a compact cycle starts. + EventCompactStart EventType = "compact_start" + // EventCompactDone is emitted when a compact cycle completes. + EventCompactDone EventType = "compact_done" + // EventCompactError is emitted when compact fails. + EventCompactError EventType = "compact_error" + // EventMicroCompactApplied is emitted when micro compact has rewritten context. + EventMicroCompactApplied EventType = "micro_compact_applied" ) diff --git a/internal/runtime/runtime.go b/internal/runtime/runtime.go index a9ea5dac..1fe8cef1 100644 --- a/internal/runtime/runtime.go +++ b/internal/runtime/runtime.go @@ -12,6 +12,7 @@ import ( "neo-code/internal/config" agentcontext "neo-code/internal/context" + contextcompact "neo-code/internal/context/compact" "neo-code/internal/provider" "neo-code/internal/tools" ) @@ -30,6 +31,7 @@ const ( type Runtime interface { Run(ctx context.Context, input UserInput) error + Compact(ctx context.Context, input CompactInput) (CompactResult, error) CancelActiveRun() bool Events() <-chan RuntimeEvent ListSessions(ctx context.Context) ([]SessionSummary, error) @@ -42,6 +44,36 @@ type UserInput struct { Content string } +type CompactInput struct { + SessionID string + RunID string +} + +type CompactResult struct { + Applied bool + BeforeChars int + AfterChars int + SavedRatio float64 + TriggerMode string + TranscriptID string + TranscriptPath string +} + +type CompactDonePayload struct { + Applied bool `json:"applied"` + BeforeChars int `json:"before_chars"` + AfterChars int `json:"after_chars"` + SavedRatio float64 `json:"saved_ratio"` + TriggerMode string `json:"trigger_mode"` + TranscriptID string `json:"transcript_id"` + TranscriptPath string `json:"transcript_path"` +} + +type CompactErrorPayload struct { + TriggerMode string `json:"trigger_mode"` + Message string `json:"message"` +} + type ProviderFactory interface { Build(ctx context.Context, cfg config.ResolvedProviderConfig) (provider.Provider, error) } @@ -52,11 +84,12 @@ type Service struct { toolManager tools.Manager // 工具管理器，统一工具 schema 暴露与执行入口。 providerFactory ProviderFactory // Provider 工厂接口，根据配置动态创建具体的 provider 实例。 contextBuilder agentcontext.Builder // 上下文构建器，负责组装 system prompt 与本轮发给模型的消息上下文。 - events chan RuntimeEvent // 事件通道，Runtime 在运行过程中产生的事件都通过该通道发送给 TUI 层消费和展示。 - runMu sync.Mutex // 运行互斥锁，保证同一时间只有一个 Run 在执行。 - activeRunToken uint64 // 当前活跃运行的令牌标识，用于标记正在执行的 Run 实例。 - nextRunToken uint64 // 下一个运行令牌的递增计数器，用于区分不同 Run 的生命周期。 - activeRunCancel context.CancelFunc // 当前活跃 Run 的取消函数。 + compactRunner contextcompact.Runner + events chan RuntimeEvent // 事件通道，Runtime 在运行过程中产生的事件都通过该通道发送给 TUI 层消费和展示。 + runMu sync.Mutex // 运行互斥锁，保证同一时间只有一个 Run 在执行。 + activeRunToken uint64 // 当前活跃运行的令牌标识，用于标记正在执行的 Run 实例。 + nextRunToken uint64 // 下一个运行令牌的递增计数器，用于区分不同 Run 的生命周期。 + activeRunCancel context.CancelFunc // 当前活跃 Run 的取消函数。 } func NewWithFactory( @@ -82,6 +115,7 @@ func NewWithFactory( toolManager: toolManager, providerFactory: providerFactory, contextBuilder: contextBuilder, + compactRunner: contextcompact.NewRunner(), events: make(chan RuntimeEvent, 128), } } @@ -131,6 +165,11 @@ func (s *Service) Run(ctx context.Context, input UserInput) error { return err } + session, _, err = s.runCompactForSession(ctx, input.RunID, session, cfg, contextcompact.ModeMicro, false) + if err != nil { + return s.handleRunError(ctx, input.RunID, session.ID, err) + } + builtContext, err := s.contextBuilder.Build(ctx, agentcontext.BuildInput{ Messages: session.Messages, Metadata: agentcontext.Metadata{ @@ -246,6 +285,36 @@ func (s *Service) Run(ctx context.Context, input UserInput) error { } } +func (s *Service) Compact(ctx context.Context, input CompactInput) (CompactResult, error) { + if err := ctx.Err(); err != nil { + return CompactResult{}, err + } + if strings.TrimSpace(input.SessionID) == "" { + return CompactResult{}, errors.New("runtime: compact session_id is empty") + } + + cfg := s.configManager.Get() + session, err := s.sessionStore.Load(ctx, input.SessionID) + if err != nil { + return CompactResult{}, err + } + + session, result, err := s.runCompactForSession(ctx, input.RunID, session, cfg, contextcompact.ModeManual, true) + if err != nil { + return CompactResult{}, err + } + + return CompactResult{ + Applied: result.Applied, + BeforeChars: result.Metrics.BeforeChars, + AfterChars: result.Metrics.AfterChars, + SavedRatio: result.Metrics.SavedRatio, + TriggerMode: result.Metrics.TriggerMode, + TranscriptID: result.TranscriptID, + TranscriptPath: result.TranscriptPath, + }, nil +} + func (s *Service) CancelActiveRun() bool { s.runMu.Lock() cancel := s.activeRunCancel @@ -364,6 +433,72 @@ func (s *Service) handleRunError(ctx context.Context, runID string, sessionID st return err } +func (s *Service) runCompactForSession( + ctx context.Context, + runID string, + session Session, + cfg config.Config, + mode contextcompact.Mode, + failOnError bool, +) (Session, contextcompact.Result, error) { + if s.compactRunner == nil { + return session, contextcompact.Result{}, nil + } + + originalMessages := append([]provider.Message(nil), session.Messages...) + s.emit(ctx, EventCompactStart, runID, session.ID, string(mode)) + + result, err := s.compactRunner.Run(ctx, contextcompact.Input{ + Mode: mode, + SessionID: session.ID, + Workdir: cfg.Workdir, + Messages: session.Messages, + Config: cfg.Context.Compact, + }) + if err != nil { + s.emit(ctx, EventCompactError, runID, session.ID, CompactErrorPayload{ + TriggerMode: string(mode), + Message: err.Error(), + }) + if failOnError { + return session, contextcompact.Result{}, err + } + return session, contextcompact.Result{}, nil + } + + if result.Applied { + session.Messages = append([]provider.Message(nil), result.Messages...) + session.UpdatedAt = time.Now() + if err := s.sessionStore.Save(ctx, &session); err != nil { + s.emit(ctx, EventCompactError, runID, session.ID, CompactErrorPayload{ + TriggerMode: string(mode), + Message: err.Error(), + }) + session.Messages = originalMessages + if failOnError { + return session, contextcompact.Result{}, err + } + return session, contextcompact.Result{}, nil + } + } + + donePayload := CompactDonePayload{ + Applied: result.Applied, + BeforeChars: result.Metrics.BeforeChars, + AfterChars: result.Metrics.AfterChars, + SavedRatio: result.Metrics.SavedRatio, + TriggerMode: result.Metrics.TriggerMode, + TranscriptID: result.TranscriptID, + TranscriptPath: result.TranscriptPath, + } + s.emit(ctx, EventCompactDone, runID, session.ID, donePayload) + if mode == contextcompact.ModeMicro && result.Applied { + s.emit(ctx, EventMicroCompactApplied, runID, session.ID, donePayload) + } + + return session, result, nil +} + // isRetryableProviderError 判断 error 是否为可重试的 ProviderError。 func isRetryableProviderError(err error) bool { var pErr *provider.ProviderError diff --git a/internal/runtime/runtime_test.go b/internal/runtime/runtime_test.go index f75a4cf8..9ceba70e 100644 --- a/internal/runtime/runtime_test.go +++ b/internal/runtime/runtime_test.go @@ -11,6 +11,7 @@ import ( "neo-code/internal/config" agentcontext "neo-code/internal/context" + contextcompact "neo-code/internal/context/compact" "neo-code/internal/provider" "neo-code/internal/provider/builtin" "neo-code/internal/tools" @@ -203,6 +204,23 @@ type stubToolManager struct { lastInput tools.ToolCallInput } +type stubCompactRunner struct { + runFn func(ctx context.Context, input contextcompact.Input) (contextcompact.Result, error) + calls []contextcompact.Input + result contextcompact.Result + err error +} + +func (r *stubCompactRunner) Run(ctx context.Context, input contextcompact.Input) (contextcompact.Result, error) { + cloned := input + cloned.Messages = append([]provider.Message(nil), input.Messages...) + r.calls = append(r.calls, cloned) + if r.runFn != nil { + return r.runFn(ctx, input) + } + return r.result, r.err +} + func (m *stubToolManager) ListAvailableSpecs(ctx context.Context, input tools.SpecListInput) ([]provider.ToolSpec, error) { m.listCalls++ if err := ctx.Err(); err != nil { @@ -1231,6 +1249,204 @@ func TestServiceConstructorsAndDelegates(t *testing.T) { } } +func TestServiceRunAppliesMicroCompactBeforeContextBuild(t *testing.T) { + manager := newRuntimeConfigManager(t) + if err := manager.Update(context.Background(), func(cfg *config.Config) error { + cfg.Context.Compact.MicroEnabled = true + return nil + }); err != nil { + t.Fatalf("enable micro compact: %v", err) + } + + store := newMemoryStore() + session := newSession("existing") + session.ID = "session-micro" + session.Messages = []provider.Message{ + {Role: provider.RoleUser, Content: "history"}, + } + store.sessions[session.ID] = cloneSession(session) + + registry := tools.NewRegistry() + registry.Register(&stubTool{name: "filesystem_read_file", content: "ok"}) + + builder := &stubContextBuilder{} + scripted := &scriptedProvider{ + responses: []provider.ChatResponse{ + { + Message: provider.Message{Role: provider.RoleAssistant, Content: "done"}, + }, + }, + } + service := NewWithFactory(manager, registry, store, &scriptedProviderFactory{provider: scripted}, builder) + compactRunner := &stubCompactRunner{ + runFn: func(ctx context.Context, input contextcompact.Input) (contextcompact.Result, error) { + next := append([]provider.Message(nil), input.Messages...) + next = append(next, provider.Message{Role: provider.RoleAssistant, Content: "micro-summary"}) + return contextcompact.Result{ + Messages: next, + Applied: true, + Metrics: contextcompact.Metrics{ + BeforeChars: 10, + AfterChars: 6, + SavedRatio: 0.4, + TriggerMode: string(contextcompact.ModeMicro), + }, + TranscriptID: "transcript_1_session-micro", + TranscriptPath: "/tmp/transcript.jsonl", + }, nil + }, + } + service.compactRunner = compactRunner + + if err := service.Run(context.Background(), UserInput{ + SessionID: session.ID, + RunID: "run-micro", + Content: "latest", + }); err != nil { + t.Fatalf("Run() error = %v", err) + } + + if len(compactRunner.calls) == 0 || compactRunner.calls[0].Mode != contextcompact.ModeMicro { + t.Fatalf("expected micro compact call, got %+v", compactRunner.calls) + } + if len(builder.lastInput.Messages) == 0 || builder.lastInput.Messages[len(builder.lastInput.Messages)-1].Content != "micro-summary" { + t.Fatalf("expected compacted messages to be passed into context builder, got %+v", builder.lastInput.Messages) + } + + events := collectRuntimeEvents(service.Events()) + assertEventSequence(t, events, []EventType{EventCompactStart, EventCompactDone, EventMicroCompactApplied, EventAgentDone}) + donePayload, ok := findEventPayload[CompactDonePayload](events, EventCompactDone) + if !ok || !donePayload.Applied || donePayload.TriggerMode != string(contextcompact.ModeMicro) { + t.Fatalf("expected compact_done payload with applied micro mode, got %+v", donePayload) + } +} + +func TestServiceRunMicroCompactFailureDoesNotBlockMainLoop(t *testing.T) { + manager := newRuntimeConfigManager(t) + if err := manager.Update(context.Background(), func(cfg *config.Config) error { + cfg.Context.Compact.MicroEnabled = true + return nil + }); err != nil { + t.Fatalf("enable micro compact: %v", err) + } + + store := newMemoryStore() + registry := tools.NewRegistry() + registry.Register(&stubTool{name: "filesystem_read_file", content: "ok"}) + + scripted := &scriptedProvider{ + responses: []provider.ChatResponse{ + { + Message: provider.Message{Role: provider.RoleAssistant, Content: "still works"}, + }, + }, + } + service := NewWithFactory(manager, registry, store, &scriptedProviderFactory{provider: scripted}, nil) + service.compactRunner = &stubCompactRunner{err: errors.New("compact failed")} + + if err := service.Run(context.Background(), UserInput{RunID: "run-micro-fail", Content: "hello"}); err != nil { + t.Fatalf("expected run to continue when micro compact fails, got %v", err) + } + + events := collectRuntimeEvents(service.Events()) + assertEventSequence(t, events, []EventType{EventCompactStart, EventCompactError, EventAgentDone}) + assertNoEventType(t, events, EventError) +} + +func TestServiceCompactManualAppliesAndPersists(t *testing.T) { + manager := newRuntimeConfigManager(t) + store := newMemoryStore() + session := newSession("manual") + session.ID = "session-manual" + session.Messages = []provider.Message{ + {Role: provider.RoleUser, Content: "before"}, + } + store.sessions[session.ID] = cloneSession(session) + + registry := tools.NewRegistry() + registry.Register(&stubTool{name: "filesystem_read_file", content: "ok"}) + + service := NewWithFactory(manager, registry, store, &scriptedProviderFactory{provider: &scriptedProvider{}}, nil) + service.compactRunner = &stubCompactRunner{ + result: contextcompact.Result{ + Messages: []provider.Message{ + {Role: provider.RoleAssistant, Content: "[compact_summary]\\ndone:\\n- ok\\n\\nin_progress:\\n- continue"}, + {Role: provider.RoleAssistant, Content: "latest"}, + }, + Applied: true, + Metrics: contextcompact.Metrics{ + BeforeChars: 80, + AfterChars: 30, + SavedRatio: 0.625, + TriggerMode: string(contextcompact.ModeManual), + }, + TranscriptID: "transcript_manual", + TranscriptPath: "/tmp/manual.jsonl", + }, + } + + result, err := service.Compact(context.Background(), CompactInput{ + SessionID: session.ID, + RunID: "run-manual", + }) + if err != nil { + t.Fatalf("Compact() error = %v", err) + } + if !result.Applied || result.BeforeChars != 80 || result.AfterChars != 30 { + t.Fatalf("unexpected compact result: %+v", result) + } + + saved, err := store.Load(context.Background(), session.ID) + if err != nil { + t.Fatalf("load compacted session: %v", err) + } + if len(saved.Messages) != 2 || !strings.Contains(saved.Messages[0].Content, "compact_summary") { + t.Fatalf("expected persisted compacted messages, got %+v", saved.Messages) + } + + events := collectRuntimeEvents(service.Events()) + assertEventSequence(t, events, []EventType{EventCompactStart, EventCompactDone}) + donePayload, ok := findEventPayload[CompactDonePayload](events, EventCompactDone) + if !ok || !donePayload.Applied || donePayload.TriggerMode != string(contextcompact.ModeManual) { + t.Fatalf("expected compact_done payload with applied manual mode, got %+v", donePayload) + } +} + +func TestServiceCompactManualFailureReturnsError(t *testing.T) { + manager := newRuntimeConfigManager(t) + store := newMemoryStore() + session := newSession("manual-fail") + session.ID = "session-manual-fail" + session.Messages = []provider.Message{{Role: provider.RoleUser, Content: "before"}} + store.sessions[session.ID] = cloneSession(session) + + registry := tools.NewRegistry() + registry.Register(&stubTool{name: "filesystem_read_file", content: "ok"}) + + service := NewWithFactory(manager, registry, store, &scriptedProviderFactory{provider: &scriptedProvider{}}, nil) + service.compactRunner = &stubCompactRunner{err: errors.New("manual compact failed")} + + _, err := service.Compact(context.Background(), CompactInput{ + SessionID: session.ID, + RunID: "run-manual-fail", + }) + if err == nil || !strings.Contains(err.Error(), "manual compact failed") { + t.Fatalf("expected compact failure, got %v", err) + } + + saved, err := store.Load(context.Background(), session.ID) + if err != nil { + t.Fatalf("load original session: %v", err) + } + if len(saved.Messages) != 1 || saved.Messages[0].Content != "before" { + t.Fatalf("expected original session untouched, got %+v", saved.Messages) + } + + events := collectRuntimeEvents(service.Events()) + assertEventSequence(t, events, []EventType{EventCompactStart, EventCompactError}) + assertNoEventType(t, events, EventCompactDone) +} + func newRuntimeConfigManager(t *testing.T) *config.Manager { t.Helper() restoreRuntimeEnv(t, config.OpenAIDefaultAPIKeyEnv) @@ -1312,6 +1528,21 @@ func assertNoEventType(t *testing.T, events []RuntimeEvent, unexpected EventType } } +func findEventPayload[T any](events []RuntimeEvent, kind EventType) (T, bool) { + var zero T + for _, event := range events { + if event.Type != kind { + continue + } + payload, ok := event.Payload.(T) + if !ok { + return zero, false + } + return payload, true + } + return zero, false +} + func assertEventsRunID(t *testing.T, events []RuntimeEvent, runID string) { t.Helper() for _, event := range events { diff --git a/internal/tui/commands.go b/internal/tui/commands.go index 18dfea81..094d4dc0 100644 --- a/internal/tui/commands.go +++ b/internal/tui/commands.go @@ -17,6 +17,7 @@ const ( slashCommandHelp = "/help" slashCommandExit = "/exit" slashCommandClear = "/clear" + slashCommandCompact = "/compact" slashCommandStatus = "/status" slashCommandProvider = "/provider" slashCommandModelPick = "/model" @@ -24,6 +25,7 @@ const ( slashUsageHelp = "/help" slashUsageExit = "/exit" slashUsageClear = "/clear" + slashUsageCompact = "/compact" slashUsageStatus = "/status" slashUsageProvider = "/provider" slashUsageModel = "/model" @@ -56,6 +58,7 @@ const ( statusApplyingCommand = "Applying local command" statusRunningCommand = "Running command" statusCommandDone = "Command finished" + statusCompacting = "Compacting context" statusChooseProvider = "Choose a provider" statusChooseModel = "Choose a model" @@ -102,6 +105,7 @@ type statusSnapshot struct { var builtinSlashCommands = []slashCommand{ {Usage: slashUsageHelp, Description: "Show slash command help"}, {Usage: slashUsageClear, Description: "Clear the current draft transcript"}, + {Usage: slashUsageCompact, Description: "Compact the current session context"}, {Usage: slashUsageStatus, Description: "Show current session and agent status"}, {Usage: slashUsageProvider, Description: "Open the interactive provider picker"}, {Usage: slashUsageModel, Description: "Open the interactive model picker"}, diff --git a/internal/tui/commands_test.go b/internal/tui/commands_test.go index 903bdaa9..32c0747d 100644 --- a/internal/tui/commands_test.go +++ b/internal/tui/commands_test.go @@ -25,6 +25,7 @@ func TestExecuteLocalCommand(t *testing.T) { for _, want := range []string{ slashUsageHelp, slashUsageClear, + slashUsageCompact, slashUsageStatus, slashUsageProvider, slashUsageModel, diff --git a/internal/tui/update.go b/internal/tui/update.go index 58747261..a37e44d6 100644 --- a/internal/tui/update.go +++ b/internal/tui/update.go @@ -21,6 +21,9 @@ import ( type RuntimeMsg struct{ Event agentruntime.RuntimeEvent } type RuntimeClosedMsg struct{} type runFinishedMsg struct{ err error } +type compactFinishedMsg struct { + err error +} type localCommandResultMsg struct { notice string err error @@ -78,6 +81,26 @@ func (a App) Update(msg tea.Msg) (tea.Model, tea.Cmd) { a.syncActiveSessionTitle() a.rebuildTranscript() return a, tea.Batch(cmds...) + case compactFinishedMsg: + // compact feedback should primarily come from runtime compact_done/error events. + if typed.err != nil && strings.TrimSpace(a.state.ExecutionError) == "" { + a.state.ExecutionError = typed.err.Error() + a.state.StatusText = typed.err.Error() + } + if err := a.refreshSessions(); err != nil { + a.state.ExecutionError = err.Error() + a.state.StatusText = err.Error() + a.appendInlineMessage(roleError, err.Error()) + } + if err := a.refreshMessages(); err != nil && strings.TrimSpace(a.state.ActiveSessionID) != "" { + a.state.ExecutionError = err.Error() + a.state.StatusText = err.Error() + a.appendInlineMessage(roleError, err.Error()) + } + a.syncActiveSessionTitle() + a.rebuildTranscript() + a.transcript.GotoBottom() + return a, tea.Batch(cmds...) case localCommandResultMsg: if typed.err != nil { a.state.ExecutionError = typed.err.Error() @@ -474,6 +497,47 @@ func (a *App) handleRuntimeEvent(event agentruntime.RuntimeEvent) { a.state.StatusText = payload a.appendInlineMessage(roleError, payload) } + case agentruntime.EventCompactDone: + payload, ok := event.Payload.(agentruntime.CompactDonePayload) + if !ok { + return + } + // Skip passive micro checks that did not rewrite context to avoid transcript noise. + if payload.TriggerMode == "micro" && !payload.Applied { + return + } + a.state.ExecutionError = "" + a.state.StatusText = fmt.Sprintf( + "Compact(%s) saved %.1f%% context", + payload.TriggerMode, + payload.SavedRatio*100, + ) + a.appendInlineMessage( + roleSystem, + fmt.Sprintf( + "[System] Compact(%s) %s (before=%d, after=%d, saved=%.1f%%, transcript=%s)", + payload.TriggerMode, + map[bool]string{true: "applied", false: "checked"}[payload.Applied], + payload.BeforeChars, + payload.AfterChars, + payload.SavedRatio*100, + payload.TranscriptPath, + ), + ) + case agentruntime.EventCompactError: + payload, ok := event.Payload.(agentruntime.CompactErrorPayload) + if !ok { + return + } + if payload.TriggerMode == "micro" { + // Micro compact failure is non-blocking for the main loop, keep it as a lightweight notice. + a.appendInlineMessage(roleEvent, fmt.Sprintf("Compact(micro) skipped: %s", payload.Message)) + return + } + message := fmt.Sprintf("Compact(%s) failed: %s", payload.TriggerMode, payload.Message) + a.state.ExecutionError = message + a.state.StatusText = message + a.appendInlineMessage(roleError, message) } } @@ -649,7 +713,7 @@ func (a *App) rebuildTranscript() { } func (a *App) handleImmediateSlashCommand(input string) (bool, tea.Cmd) { - command, _ := splitFirstWord(strings.ToLower(strings.TrimSpace(input))) + command, rest := splitFirstWord(strings.ToLower(strings.TrimSpace(input))) switch command { case slashCommandExit: return true, tea.Quit @@ -657,6 +721,18 @@ func (a *App) handleImmediateSlashCommand(input string) (bool, tea.Cmd) { a.startDraftSession() a.state.StatusText = "[System] Cleared current draft/history." return true, nil + case slashCommandCompact: + if strings.TrimSpace(rest) != "" { + errText := fmt.Sprintf("usage: %s", slashUsageCompact) + a.state.ExecutionError = errText + a.state.StatusText = errText + a.appendInlineMessage(roleError, errText) + a.rebuildTranscript() + return true, nil + } + a.state.StatusText = statusCompacting + a.state.ExecutionError = "" + return true, runCompact(a.runtime, a.state.ActiveSessionID) default: return false, nil } @@ -717,3 +793,10 @@ func runAgent(runtime agentruntime.Runtime, sessionID string, content string) te return runFinishedMsg{err: err} } } + +func runCompact(runtime agentruntime.Runtime, sessionID string) tea.Cmd { + return func() tea.Msg { + _, err := runtime.Compact(context.Background(), agentruntime.CompactInput{SessionID: sessionID}) + return compactFinishedMsg{err: err} + } +} diff --git a/internal/tui/update_test.go b/internal/tui/update_test.go index f9d7e061..5220507a 100644 --- a/internal/tui/update_test.go +++ b/internal/tui/update_test.go @@ -19,15 +19,18 @@ import ( ) type stubRuntime struct { - runInputs []agentruntime.UserInput - events chan agentruntime.RuntimeEvent - sessions []agentruntime.SessionSummary - loads map[string]agentruntime.Session - runErr error - listErr error - loadErr error - cancelCalls int - cancelResult bool + runInputs []agentruntime.UserInput + compactInputs []agentruntime.CompactInput + events chan agentruntime.RuntimeEvent + sessions []agentruntime.SessionSummary + loads map[string]agentruntime.Session + runErr error + compactErr error + compactResult agentruntime.CompactResult + listErr error + loadErr error + cancelCalls int + cancelResult bool } func newStubRuntime() *stubRuntime { @@ -42,6 +45,11 @@ func (r *stubRuntime) Run(ctx context.Context, input agentruntime.UserInput) err return r.runErr } +func (r *stubRuntime) Compact(ctx context.Context, input agentruntime.CompactInput) (agentruntime.CompactResult, error) { + r.compactInputs = append(r.compactInputs, input) + return r.compactResult, r.compactErr +} + func (r *stubRuntime) Events() <-chan agentruntime.RuntimeEvent { return r.events } @@ -1133,6 +1141,93 @@ func TestAppHandleRuntimeEventAdditionalBranches(t *testing.T) { } }, }, + { + name: "compact done event emits lightweight notice", + event: agentruntime.RuntimeEvent{ + Type: agentruntime.EventCompactDone, + SessionID: "s1", + Payload: agentruntime.CompactDonePayload{ + Applied: true, + BeforeChars: 100, + AfterChars: 60, + SavedRatio: 0.4, + TriggerMode: "manual", + TranscriptPath: "/tmp/t.jsonl", + }, + }, + assert: func(t *testing.T, app App) { + t.Helper() + if !strings.Contains(app.state.StatusText, "Compact(manual)") { + t.Fatalf("expected compact status text, got %q", app.state.StatusText) + } + if len(app.activeMessages) == 0 || !strings.Contains(app.activeMessages[len(app.activeMessages)-1].Content, "Compact(manual)") { + t.Fatalf("expected compact inline notice, got %+v", app.activeMessages) + } + }, + }, + { + name: "compact done ignores passive micro checks", + setup: func(app *App) { + app.state.StatusText = "unchanged" + }, + event: agentruntime.RuntimeEvent{ + Type: agentruntime.EventCompactDone, + SessionID: "s1", + Payload: agentruntime.CompactDonePayload{ + Applied: false, + TriggerMode: "micro", + }, + }, + assert: func(t *testing.T, app App) { + t.Helper() + if app.state.StatusText != "unchanged" { + t.Fatalf("expected status unchanged for passive micro compact, got %q", app.state.StatusText) + } + }, + }, + { + name: "compact error event is surfaced", + event: agentruntime.RuntimeEvent{ + Type: agentruntime.EventCompactError, + SessionID: "s1", + Payload: agentruntime.CompactErrorPayload{ + TriggerMode: "manual", + Message: "disk full", + }, + }, + assert: func(t *testing.T, app App) { + t.Helper() + if !strings.Contains(app.state.ExecutionError, "Compact(manual) failed: disk full") { + t.Fatalf("expected compact error in state, got %q", app.state.ExecutionError) + } + }, + }, + { + name: "micro compact error remains non-blocking notice", + setup: func(app *App) { + app.state.StatusText = statusThinking + }, + event: agentruntime.RuntimeEvent{ + Type: agentruntime.EventCompactError, + SessionID: "s1", + Payload: agentruntime.CompactErrorPayload{ + TriggerMode: "micro", + Message: "tmp write failed", + }, + }, + assert: func(t *testing.T, app App) { + t.Helper() + if app.state.ExecutionError != "" { + t.Fatalf("expected no execution error for micro compact failure, got %q", app.state.ExecutionError) + } + if app.state.StatusText != statusThinking { + t.Fatalf("expected status unchanged for micro compact failure, got %q", app.state.StatusText) + } + if len(app.activeMessages) == 0 || !strings.Contains(app.activeMessages[len(app.activeMessages)-1].Content, "Compact(micro) skipped") { + t.Fatalf("expected lightweight micro compact notice, got %+v", app.activeMessages) + } + }, + }, } for _, tt := range tests { @@ -1205,6 +1300,40 @@ func TestImmediateSlashCommandsAndLayoutBranches(t *testing.T) { t.Fatalf("expected /clear to reset draft state") } + runtime.compactResult = agentruntime.CompactResult{ + Applied: true, + BeforeChars: 100, + AfterChars: 40, + SavedRatio: 0.6, + TranscriptPath: "/tmp/transcript.jsonl", + } + app.state.ActiveSessionID = "session-compact" + handled, cmd = app.handleImmediateSlashCommand("/compact") + if !handled || cmd == nil { + t.Fatalf("expected /compact to trigger compact cmd") + } + if app.state.StatusText != statusCompacting { + t.Fatalf("expected compact status %q, got %q", statusCompacting, app.state.StatusText) + } + msgs := collectTeaMessages(cmd) + if len(msgs) != 1 { + t.Fatalf("expected one compact command message, got %d", len(msgs)) + } + if _, ok := msgs[0].(compactFinishedMsg); !ok { + t.Fatalf("expected compact finished msg, got %T", msgs[0]) + } + if len(runtime.compactInputs) != 1 || runtime.compactInputs[0].SessionID != "session-compact" { + t.Fatalf("expected runtime compact call with active session, got %+v", runtime.compactInputs) + } + + handled, cmd = app.handleImmediateSlashCommand("/compact now") + if !handled || cmd != nil { + t.Fatalf("expected /compact with args to be handled locally with usage error") + } + if !strings.Contains(app.state.StatusText, "usage: /compact") { + t.Fatalf("expected /compact usage hint, got %q", app.state.StatusText) + } + handled, cmd = app.handleImmediateSlashCommand("/exit") if !handled || cmd == nil { t.Fatalf("expected /exit to return a quit cmd") From 80d000b6bf01420d7efedd8f36649e85cef19b2d Mon Sep 17 00:00:00 2001 From: Yumiue <229866007@qq.com> Date: Thu, 2 Apr 2026 22:30:13 +0800 Subject: [PATCH 02/23] =?UTF-8?q?=EF=BB=BFfeat(compact):=20=E5=AE=8C?= =?UTF-8?q?=E6=88=90=20M1=20=E6=94=B6=E5=8F=A3=E5=B9=B6=E4=BF=AE=E5=A4=8D?= =?UTF-8?q?=E5=85=B3=E9=94=AE=E9=A3=8E=E9=99=A9?= MIME-Version: 1.0 Content-Type: text/plain; charset=UTF-8 Content-Transfer-Encoding: 8bit - 默认启用 micro_compact，新增 manual full_replace 策略 - 统一 compact 流程并补齐 transcript 持久化与唯一命名 - 修复 micro_enabled=false 持久化丢失问题，兼容缺失字段回退默认值 - 调整 TUI compact 轻量可观测，避免 micro 失败污染主流程错误态 - 同步配置/事件/持久化相关文档并补齐回归测试 测试： - go test ./... --- .gitignore | 1 + README.md | 3 + docs/config-management-detail-design.md | 11 + docs/runtime-provider-event-flow.md | 31 +- docs/session-persistence-design.md | 6 + docs/tools-and-tui-integration.md | 2 + internal/config/config_test.go | 178 +++++++++ internal/config/loader.go | 72 +++- internal/config/model.go | 123 +++++- internal/context/compact/runner.go | 500 ++++++++++++++++++++++++ internal/context/compact/runner_test.go | 337 ++++++++++++++++ internal/runtime/events.go | 8 + internal/runtime/runtime.go | 145 ++++++- internal/runtime/runtime_test.go | 231 +++++++++++ internal/tui/commands.go | 4 + internal/tui/commands_test.go | 1 + internal/tui/update.go | 88 ++++- internal/tui/update_test.go | 147 ++++++- 18 files changed, 1853 insertions(+), 35 deletions(-) create mode 100644 internal/context/compact/runner.go create mode 100644 internal/context/compact/runner_test.go diff --git a/.gitignore b/.gitignore index df09beca..8c4a7761 100644 --- a/.gitignore +++ b/.gitignore @@ -31,6 +31,7 @@ config.local.yaml # Local data data/ .cache/ +.neocode/projects/**/.transcripts/ # Editor/IDE .idea/ diff --git a/README.md b/README.md index 317e8c20..6c6d15ea 100644 --- a/README.md +++ b/README.md @@ -68,6 +68,9 @@ go run ./cmd/neocode - `/provider` — 切换内建模型提供商 - `/model` — 切换模型 +- `/compact` — 手动压缩当前会话上下文（保留可追溯 transcript） + +上下文压缩默认开启每轮 `micro_compact`，手动 `/compact` 支持 `keep_recent` 与 `full_replace` 两种策略。 ## 架构概览 diff --git a/docs/config-management-detail-design.md b/docs/config-management-detail-design.md index 8975eb90..9e79bb52 100644 --- a/docs/config-management-detail-design.md +++ b/docs/config-management-detail-design.md @@ -8,6 +8,17 @@ ## 核心类型 - `Config`：顶层应用配置，运行时包含 Provider 列表、当前选中 Provider、当前模型、工作目录、Shell 和循环限制等信息 +- `Config.Context.Compact`：上下文压缩配置，包含： + - `micro_enabled` + - `tool_result_keep_recent` + - `tool_result_placeholder_min_chars` + - `manual_strategy`（M1 支持 `keep_recent` / `full_replace`） + - `manual_keep_recent_spans` + - `max_summary_chars` + +默认策略说明： +- `micro_enabled` 默认 `true`，可通过配置显式关闭 +- `manual_strategy` 默认 `keep_recent` - `ProviderConfig`：单个 Provider 的内建定义，包括 Base URL、默认模型、实例级模型列表和 API Key 环境变量名 - `Manager`：使用 `sync.RWMutex` 保护的配置访问器与修改器 - `Loader`：对 YAML 文件和托管 `.env` 文件的文件系统封装 diff --git a/docs/runtime-provider-event-flow.md b/docs/runtime-provider-event-flow.md index e9e007e2..83a15e37 100644 --- a/docs/runtime-provider-event-flow.md +++ b/docs/runtime-provider-event-flow.md @@ -9,18 +9,23 @@ - `tool_start` - `tool_result` - `error` +- `compact_start` +- `compact_done` +- `compact_error` +- `micro_compact_applied` ## ReAct 主循环 1. 加载目标会话或创建新会话。 2. 追加最新的用户消息。 3. 读取最新配置快照。 -4. 解析当前 provider 配置并构建 provider 实例。 -5. 调用 `context.Builder` 生成本轮请求使用的 `system prompt` 和消息上下文。 -6. 调用 `Provider.Chat`，并把流式事件桥接给 TUI。 -7. 保存 assistant 完整回复。 -8. 执行返回的工具调用，并保存每一个工具结果。 -9. 如果仍需继续推理，则进入下一轮；否则结束。 +4. 每轮请求前先尝试执行 `micro_compact`（失败时降级跳过，不阻断主链路）。 +5. 解析当前 provider 配置并构建 provider 实例。 +6. 调用 `context.Builder` 生成本轮请求使用的 `system prompt` 和消息上下文。 +7. 调用 `Provider.Chat`，并把流式事件桥接给 TUI。 +8. 保存 assistant 完整回复。 +9. 执行返回的工具调用，并保存每一个工具结果。 +10. 如果仍需继续推理，则进入下一轮；否则结束。 ### Context Builder 输入与职责 @@ -59,6 +64,20 @@ - runtime 将其转换成 `RuntimeEvent` - TUI 使用 Bubble Tea `Cmd` 监听事件，并在处理完成后继续订阅 +### Compact 事件 + +- `compact_start`：开始执行 compact（payload 为触发模式，`micro` 或 `manual`） +- `compact_done`：compact 完成，payload 包含： + - `applied` + - `before_chars` + - `after_chars` + - `saved_ratio` + - `trigger_mode` + - `transcript_id` + - `transcript_path` +- `compact_error`：compact 失败，payload 包含触发模式和错误信息 +- `micro_compact_applied`：仅当 `micro_compact` 实际改写了消息时触发 + ## 持久化时机 - 用户消息提交后保存 diff --git a/docs/session-persistence-design.md b/docs/session-persistence-design.md index 0c5d1887..4f5d1cde 100644 --- a/docs/session-persistence-design.md +++ b/docs/session-persistence-design.md @@ -2,14 +2,18 @@ ## 存储策略 NeoCode 在 MVP 阶段使用 JSON 文件持久化 Session，以保持本地优先、易于调试和跨平台可移植。 +除会话主存储外，compact 会在压缩前额外写入 transcript 留痕，确保“移出活跃上下文”而非“不可恢复删除”。 + ## 数据模型 - `Session`：完整消息历史以及 `id`、`title`、`updated_at` 等元信息 - `SessionSummary`：用于侧边栏的轻量摘要结构 +- `Transcript(.jsonl)`：compact 前完整消息快照，每行一条消息，包含 `role/content/tool_calls/tool_call_id/is_error/index/timestamp` ## 加载策略 - `ListSummaries` 只读取渲染侧边栏所需的基础信息 - `Load` 仅在用户真正进入某个会话时读取完整消息历史 - `Save` 通过临时文件原子写入完整 Session +- transcript 当前不参与自动加载回放，仅用于追溯与审计 ## 命名策略 - 新会话默认展示为 `Draft` @@ -18,3 +22,5 @@ NeoCode 在 MVP 阶段使用 JSON 文件持久化 Session，以保持本地优 ## 并发约束 - SessionStore 实现必须自行保护共享访问 - 真正的保存时机由 runtime 决定，TUI 不负责直接触发磁盘写入 +- compact transcript 也通过原子落盘写入，路径为 `~/.neocode/projects//.transcripts/` +- transcript 文件在 Unix 系统默认使用 `0600` 权限，减少敏感内容泄露面 diff --git a/docs/tools-and-tui-integration.md b/docs/tools-and-tui-integration.md index b41a38a3..c092536e 100644 --- a/docs/tools-and-tui-integration.md +++ b/docs/tools-and-tui-integration.md @@ -22,6 +22,8 @@ ## TUI 集成方式 - 本地配置操作统一通过 Slash Command 完成，例如 Base URL、API Key 和模型选择 +- `/compact` 提供手动上下文压缩入口，触发 runtime 的 manual compact 流程 +- TUI 对 compact 仅做轻量可观测展示：显示 `compact_done` / `compact_error` 关键信息，避免事件刷屏 - runtime 事件以内联形式渲染到 transcript 中，而不是单独拆出控制台面板 - 工具开始和结束事件会以轻量提示插入聊天流，使交互更沉浸 diff --git a/internal/config/config_test.go b/internal/config/config_test.go index cc0d4ead..6e515365 100644 --- a/internal/config/config_test.go +++ b/internal/config/config_test.go @@ -829,6 +829,184 @@ func TestConstructorsRejectMissingDependencies(t *testing.T) { }) } +func TestCompactConfigDefaultsAndRoundTrip(t *testing.T) { + tempDir := t.TempDir() + loader := NewLoader(tempDir, testDefaultConfig()) + + cfg, err := loader.Load(context.Background()) + if err != nil { + t.Fatalf("Load() error = %v", err) + } + + compactCfg := cfg.Context.Compact + if !compactCfg.MicroEnabled { + t.Fatalf("expected micro_enabled default true") + } + if compactCfg.ToolResultKeepRecent != DefaultCompactToolResultKeepRecent { + t.Fatalf("expected tool_result_keep_recent=%d, got %d", DefaultCompactToolResultKeepRecent, compactCfg.ToolResultKeepRecent) + } + if compactCfg.ToolResultPlaceholderMinChars != DefaultCompactPlaceholderMinChars { + t.Fatalf("expected tool_result_placeholder_min_chars=%d, got %d", DefaultCompactPlaceholderMinChars, compactCfg.ToolResultPlaceholderMinChars) + } + if compactCfg.ManualStrategy != CompactManualStrategyKeepRecent { + t.Fatalf("expected manual strategy %q, got %q", CompactManualStrategyKeepRecent, compactCfg.ManualStrategy) + } + if compactCfg.ManualKeepRecentSpans != DefaultCompactManualKeepRecentSpans { + t.Fatalf("expected manual_keep_recent_spans=%d, got %d", DefaultCompactManualKeepRecentSpans, compactCfg.ManualKeepRecentSpans) + } + if compactCfg.MaxSummaryChars != DefaultCompactMaxSummaryChars { + t.Fatalf("expected max_summary_chars=%d, got %d", DefaultCompactMaxSummaryChars, compactCfg.MaxSummaryChars) + } + + if err := loader.Save(context.Background(), cfg); err != nil { + t.Fatalf("Save() error = %v", err) + } + reloaded, err := loader.Load(context.Background()) + if err != nil { + t.Fatalf("Reload() error = %v", err) + } + + if reloaded.Context.Compact.ManualStrategy != CompactManualStrategyKeepRecent { + t.Fatalf("expected manual strategy to persist, got %q", reloaded.Context.Compact.ManualStrategy) + } +} + +func TestCompactConfigMicroEnabledFalsePersistsAcrossReload(t *testing.T) { + tempDir := t.TempDir() + loader := NewLoader(tempDir, testDefaultConfig()) + + cfg, err := loader.Load(context.Background()) + if err != nil { + t.Fatalf("Load() error = %v", err) + } + cfg.Context.Compact.MicroEnabled = false + + if err := loader.Save(context.Background(), cfg); err != nil { + t.Fatalf("Save() error = %v", err) + } + reloaded, err := loader.Load(context.Background()) + if err != nil { + t.Fatalf("Reload() error = %v", err) + } + if reloaded.Context.Compact.MicroEnabled { + t.Fatalf("expected micro_enabled to persist as false after reload") + } +} + +func TestCompactConfigMissingMicroEnabledFallsBackToDefaultTrue(t *testing.T) { + tempDir := t.TempDir() + loader := NewLoader(tempDir, testDefaultConfig()) + if err := os.MkdirAll(loader.BaseDir(), 0o755); err != nil { + t.Fatalf("create base dir: %v", err) + } + + legacyConfig := []byte(strings.Join([]string{ + "selected_provider: openai", + "current_model: gpt-4.1", + "workdir: " + filepath.ToSlash(tempDir), + "shell: powershell", + "", + }, "\n")) + if err := os.WriteFile(loader.ConfigPath(), legacyConfig, 0o644); err != nil { + t.Fatalf("write legacy config: %v", err) + } + + cfg, err := loader.Load(context.Background()) + if err != nil { + t.Fatalf("Load() error = %v", err) + } + if !cfg.Context.Compact.MicroEnabled { + t.Fatalf("expected missing micro_enabled to fallback to default true") + } +} + +func TestCompactConfigValidateFailures(t *testing.T) { + tests := []struct { + name string + compact CompactConfig + expectErr string + }{ + { + name: "invalid keep recent", + compact: CompactConfig{ + ToolResultKeepRecent: 0, + ToolResultPlaceholderMinChars: 100, + ManualStrategy: CompactManualStrategyKeepRecent, + ManualKeepRecentSpans: 6, + MaxSummaryChars: 1200, + }, + expectErr: "tool_result_keep_recent", + }, + { + name: "invalid placeholder min chars", + compact: CompactConfig{ + ToolResultKeepRecent: 3, + ToolResultPlaceholderMinChars: 0, + ManualStrategy: CompactManualStrategyKeepRecent, + ManualKeepRecentSpans: 6, + MaxSummaryChars: 1200, + }, + expectErr: "tool_result_placeholder_min_chars", + }, + { + name: "invalid manual strategy", + compact: CompactConfig{ + ToolResultKeepRecent: 3, + ToolResultPlaceholderMinChars: 100, + ManualStrategy: "unsupported", + ManualKeepRecentSpans: 6, + MaxSummaryChars: 1200, + }, + expectErr: "manual_strategy", + }, + { + name: "invalid manual keep spans", + compact: CompactConfig{ + ToolResultKeepRecent: 3, + ToolResultPlaceholderMinChars: 100, + ManualStrategy: CompactManualStrategyKeepRecent, + ManualKeepRecentSpans: 0, + MaxSummaryChars: 1200, + }, + expectErr: "manual_keep_recent_spans", + }, + { + name: "invalid summary chars", + compact: CompactConfig{ + ToolResultKeepRecent: 3, + ToolResultPlaceholderMinChars: 100, + ManualStrategy: CompactManualStrategyKeepRecent, + ManualKeepRecentSpans: 6, + MaxSummaryChars: 0, + }, + expectErr: "max_summary_chars", + }, + } + + for _, tt := range tests { + tt := tt + t.Run(tt.name, func(t *testing.T) { + err := tt.compact.Validate() + if err == nil || !strings.Contains(err.Error(), tt.expectErr) { + t.Fatalf("expected error containing %q, got %v", tt.expectErr, err) + } + }) + } +} + +func TestCompactConfigValidateSupportsFullReplace(t *testing.T) { + err := (CompactConfig{ + ToolResultKeepRecent: 3, + ToolResultPlaceholderMinChars: 100, + ManualStrategy: CompactManualStrategyFullReplace, + ManualKeepRecentSpans: 6, + MaxSummaryChars: 1200, + }).Validate() + if err != nil { + t.Fatalf("expected full_replace strategy to be valid, got %v", err) + } +} + func restoreEnv(t *testing.T, key string) { t.Helper() value, ok := os.LookupEnv(key) diff --git a/internal/config/loader.go b/internal/config/loader.go index c3563828..7ee053b1 100644 --- a/internal/config/loader.go +++ b/internal/config/loader.go @@ -22,13 +22,27 @@ type Loader struct { } type persistedConfig struct { - SelectedProvider string `yaml:"selected_provider"` - CurrentModel string `yaml:"current_model"` - Workdir string `yaml:"workdir"` - Shell string `yaml:"shell"` - MaxLoops int `yaml:"max_loops,omitempty"` - ToolTimeoutSec int `yaml:"tool_timeout_sec,omitempty"` - Tools ToolsConfig `yaml:"tools,omitempty"` + SelectedProvider string `yaml:"selected_provider"` + CurrentModel string `yaml:"current_model"` + Workdir string `yaml:"workdir"` + Shell string `yaml:"shell"` + MaxLoops int `yaml:"max_loops,omitempty"` + ToolTimeoutSec int `yaml:"tool_timeout_sec,omitempty"` + Context persistedContextConfig `yaml:"context,omitempty"` + Tools ToolsConfig `yaml:"tools,omitempty"` +} + +type persistedContextConfig struct { + Compact persistedCompactConfig `yaml:"compact,omitempty"` +} + +type persistedCompactConfig struct { + MicroEnabled *bool `yaml:"micro_enabled,omitempty"` + ToolResultKeepRecent int `yaml:"tool_result_keep_recent,omitempty"` + ToolResultPlaceholderMinChars int `yaml:"tool_result_placeholder_min_chars,omitempty"` + ManualStrategy string `yaml:"manual_strategy,omitempty"` + ManualKeepRecentSpans int `yaml:"manual_keep_recent_spans,omitempty"` + MaxSummaryChars int `yaml:"max_summary_chars,omitempty"` } func NewLoader(baseDir string, defaults *Config) *Loader { @@ -84,7 +98,7 @@ func (l *Loader) Load(ctx context.Context) (*Config, error) { return nil, fmt.Errorf("config: read config file: %w", err) } - cfg, err := parseConfig(data) + cfg, err := parseConfigWithContextDefaults(data, l.defaults.Context) if err != nil { return nil, fmt.Errorf("config: parse config file: %w", err) } @@ -140,14 +154,18 @@ func defaultBaseDir() string { } func parseConfig(data []byte) (*Config, error) { + return parseConfigWithContextDefaults(data, Default().Context) +} + +func parseConfigWithContextDefaults(data []byte, contextDefaults ContextConfig) (*Config, error) { if len(bytes.TrimSpace(data)) == 0 { return &Config{}, nil } - return parseCurrentConfig(data) + return parseCurrentConfig(data, contextDefaults) } -func parseCurrentConfig(data []byte) (*Config, error) { +func parseCurrentConfig(data []byte, contextDefaults ContextConfig) (*Config, error) { var file persistedConfig if err := yaml.Unmarshal(data, &file); err != nil { return nil, err @@ -160,6 +178,7 @@ func parseCurrentConfig(data []byte) (*Config, error) { Shell: strings.TrimSpace(file.Shell), MaxLoops: file.MaxLoops, ToolTimeoutSec: file.ToolTimeoutSec, + Context: fromPersistedContextConfig(file.Context, contextDefaults), Tools: file.Tools, } @@ -174,6 +193,7 @@ func marshalPersistedConfig(snapshot Config) ([]byte, error) { Shell: snapshot.Shell, MaxLoops: snapshot.MaxLoops, ToolTimeoutSec: snapshot.ToolTimeoutSec, + Context: newPersistedContextConfig(snapshot.Context), Tools: snapshot.Tools, } @@ -187,6 +207,38 @@ func marshalPersistedConfig(snapshot Config) ([]byte, error) { return data, nil } +func newPersistedContextConfig(cfg ContextConfig) persistedContextConfig { + microEnabled := cfg.Compact.MicroEnabled + return persistedContextConfig{ + Compact: persistedCompactConfig{ + MicroEnabled: µEnabled, + ToolResultKeepRecent: cfg.Compact.ToolResultKeepRecent, + ToolResultPlaceholderMinChars: cfg.Compact.ToolResultPlaceholderMinChars, + ManualStrategy: cfg.Compact.ManualStrategy, + ManualKeepRecentSpans: cfg.Compact.ManualKeepRecentSpans, + MaxSummaryChars: cfg.Compact.MaxSummaryChars, + }, + } +} + +func fromPersistedContextConfig(file persistedContextConfig, defaults ContextConfig) ContextConfig { + out := ContextConfig{ + Compact: CompactConfig{ + ToolResultKeepRecent: file.Compact.ToolResultKeepRecent, + ToolResultPlaceholderMinChars: file.Compact.ToolResultPlaceholderMinChars, + ManualStrategy: strings.TrimSpace(file.Compact.ManualStrategy), + ManualKeepRecentSpans: file.Compact.ManualKeepRecentSpans, + MaxSummaryChars: file.Compact.MaxSummaryChars, + }, + } + if file.Compact.MicroEnabled != nil { + out.Compact.MicroEnabled = *file.Compact.MicroEnabled + } else { + out.Compact.MicroEnabled = defaults.Compact.MicroEnabled + } + return out +} + func persistedConfigDiffers(data []byte, cfg Config) (bool, error) { canonical, err := marshalPersistedConfig(cfg) if err != nil { diff --git a/internal/config/model.go b/internal/config/model.go index f8c2bd52..4341e49c 100644 --- a/internal/config/model.go +++ b/internal/config/model.go @@ -11,10 +11,19 @@ import ( ) const ( - DefaultWorkdir = "." - DefaultMaxLoops = 8 - DefaultToolTimeoutSec = 20 - DefaultWebFetchMaxResponseBytes int64 = 256 * 1024 + DefaultWorkdir = "." + DefaultMaxLoops = 8 + DefaultToolTimeoutSec = 20 + DefaultWebFetchMaxResponseBytes int64 = 256 * 1024 + DefaultCompactToolResultKeepRecent = 3 + DefaultCompactPlaceholderMinChars = 100 + DefaultCompactManualKeepRecentSpans = 6 + DefaultCompactMaxSummaryChars = 1200 +) + +const ( + CompactManualStrategyKeepRecent = "keep_recent" + CompactManualStrategyFullReplace = "full_replace" ) var defaultWebFetchSupportedContentTypes = []string{ @@ -34,6 +43,7 @@ type Config struct { Shell string `yaml:"shell"` MaxLoops int `yaml:"max_loops,omitempty"` ToolTimeoutSec int `yaml:"tool_timeout_sec,omitempty"` + Context ContextConfig `yaml:"context,omitempty"` Tools ToolsConfig `yaml:"tools,omitempty"` } @@ -54,6 +64,19 @@ type ToolsConfig struct { WebFetch WebFetchConfig `yaml:"webfetch,omitempty"` } +type ContextConfig struct { + Compact CompactConfig `yaml:"compact,omitempty"` +} + +type CompactConfig struct { + MicroEnabled bool `yaml:"micro_enabled,omitempty"` + ToolResultKeepRecent int `yaml:"tool_result_keep_recent,omitempty"` + ToolResultPlaceholderMinChars int `yaml:"tool_result_placeholder_min_chars,omitempty"` + ManualStrategy string `yaml:"manual_strategy,omitempty"` + ManualKeepRecentSpans int `yaml:"manual_keep_recent_spans,omitempty"` + MaxSummaryChars int `yaml:"max_summary_chars,omitempty"` +} + type WebFetchConfig struct { MaxResponseBytes int64 `yaml:"max_response_bytes,omitempty"` SupportedContentTypes []string `yaml:"supported_content_types,omitempty"` @@ -69,6 +92,7 @@ func Default() *Config { Shell: defaultShell(), MaxLoops: DefaultMaxLoops, ToolTimeoutSec: DefaultToolTimeoutSec, + Context: defaultContextConfig(), Tools: ToolsConfig{ WebFetch: defaultWebFetchConfig(), }, @@ -82,6 +106,7 @@ func (c *Config) Clone() Config { clone := *c clone.Providers = cloneProviders(c.Providers) + clone.Context = c.Context.Clone() clone.Tools = c.Tools.Clone() return clone } @@ -126,6 +151,7 @@ func (c *Config) ApplyDefaultsFrom(defaults Config) { if c.ToolTimeoutSec <= 0 { c.ToolTimeoutSec = defaults.ToolTimeoutSec } + c.Context.ApplyDefaults(defaults.Context) c.Tools.ApplyDefaults(defaults.Tools) c.Workdir = normalizeWorkdir(c.Workdir) @@ -189,6 +215,9 @@ func (c *Config) Validate() error { if err := c.Tools.Validate(); err != nil { return fmt.Errorf("config: tools: %w", err) } + if err := c.Context.Validate(); err != nil { + return fmt.Errorf("config: context: %w", err) + } return nil } @@ -305,12 +334,35 @@ func defaultWebFetchConfig() WebFetchConfig { } } +func defaultContextConfig() ContextConfig { + return ContextConfig{ + Compact: defaultCompactConfig(), + } +} + +func defaultCompactConfig() CompactConfig { + return CompactConfig{ + MicroEnabled: true, + ToolResultKeepRecent: DefaultCompactToolResultKeepRecent, + ToolResultPlaceholderMinChars: DefaultCompactPlaceholderMinChars, + ManualStrategy: CompactManualStrategyKeepRecent, + ManualKeepRecentSpans: DefaultCompactManualKeepRecentSpans, + MaxSummaryChars: DefaultCompactMaxSummaryChars, + } +} + func (c ToolsConfig) Clone() ToolsConfig { return ToolsConfig{ WebFetch: c.WebFetch.Clone(), } } +func (c ContextConfig) Clone() ContextConfig { + return ContextConfig{ + Compact: c.Compact.Clone(), + } +} + func (c *ToolsConfig) ApplyDefaults(defaults ToolsConfig) { if c == nil { return @@ -319,6 +371,14 @@ func (c *ToolsConfig) ApplyDefaults(defaults ToolsConfig) { c.WebFetch.ApplyDefaults(defaults.WebFetch) } +func (c *ContextConfig) ApplyDefaults(defaults ContextConfig) { + if c == nil { + return + } + + c.Compact.ApplyDefaults(defaults.Compact) +} + func (c ToolsConfig) Validate() error { if err := c.WebFetch.Validate(); err != nil { return fmt.Errorf("webfetch: %w", err) @@ -326,12 +386,23 @@ func (c ToolsConfig) Validate() error { return nil } +func (c ContextConfig) Validate() error { + if err := c.Compact.Validate(); err != nil { + return fmt.Errorf("compact: %w", err) + } + return nil +} + func (c WebFetchConfig) Clone() WebFetchConfig { clone := c clone.SupportedContentTypes = append([]string(nil), c.SupportedContentTypes...) return clone } +func (c CompactConfig) Clone() CompactConfig { + return c +} + func (c *WebFetchConfig) ApplyDefaults(defaults WebFetchConfig) { if c == nil { return @@ -343,6 +414,28 @@ func (c *WebFetchConfig) ApplyDefaults(defaults WebFetchConfig) { c.SupportedContentTypes = normalizeContentTypes(c.SupportedContentTypes, defaults.SupportedContentTypes) } +func (c *CompactConfig) ApplyDefaults(defaults CompactConfig) { + if c == nil { + return + } + + if c.ToolResultKeepRecent <= 0 { + c.ToolResultKeepRecent = defaults.ToolResultKeepRecent + } + if c.ToolResultPlaceholderMinChars <= 0 { + c.ToolResultPlaceholderMinChars = defaults.ToolResultPlaceholderMinChars + } + if strings.TrimSpace(c.ManualStrategy) == "" { + c.ManualStrategy = defaults.ManualStrategy + } + if c.ManualKeepRecentSpans <= 0 { + c.ManualKeepRecentSpans = defaults.ManualKeepRecentSpans + } + if c.MaxSummaryChars <= 0 { + c.MaxSummaryChars = defaults.MaxSummaryChars + } +} + func (c WebFetchConfig) Validate() error { if c.MaxResponseBytes <= 0 { return errors.New("max_response_bytes must be greater than 0") @@ -359,6 +452,28 @@ func (c WebFetchConfig) Validate() error { return nil } +func (c CompactConfig) Validate() error { + if c.ToolResultKeepRecent <= 0 { + return errors.New("tool_result_keep_recent must be greater than 0") + } + if c.ToolResultPlaceholderMinChars <= 0 { + return errors.New("tool_result_placeholder_min_chars must be greater than 0") + } + if c.ManualKeepRecentSpans <= 0 { + return errors.New("manual_keep_recent_spans must be greater than 0") + } + if c.MaxSummaryChars <= 0 { + return errors.New("max_summary_chars must be greater than 0") + } + + switch strings.ToLower(strings.TrimSpace(c.ManualStrategy)) { + case CompactManualStrategyKeepRecent, CompactManualStrategyFullReplace: + return nil + default: + return fmt.Errorf("manual_strategy %q is not supported", c.ManualStrategy) + } +} + func normalizeContentTypes(values []string, defaults []string) []string { source := values if len(source) == 0 { diff --git a/internal/context/compact/runner.go b/internal/context/compact/runner.go new file mode 100644 index 00000000..2702892d --- /dev/null +++ b/internal/context/compact/runner.go @@ -0,0 +1,500 @@ +package compact + +import ( + "context" + cryptorand "crypto/rand" + "crypto/sha1" + "encoding/hex" + "encoding/json" + "errors" + "fmt" + "os" + "path/filepath" + "regexp" + goruntime "runtime" + "strings" + "time" + + "neo-code/internal/config" + "neo-code/internal/provider" +) + +type Mode string + +const ( + ModeMicro Mode = "micro" + ModeManual Mode = "manual" +) + +type ErrorMode string + +const ( + ErrorModeNone ErrorMode = "none" +) + +type Input struct { + Mode Mode + SessionID string + Workdir string + Messages []provider.Message + Config config.CompactConfig +} + +type Metrics struct { + BeforeChars int `json:"before_chars"` + AfterChars int `json:"after_chars"` + SavedRatio float64 `json:"saved_ratio"` + TriggerMode string `json:"trigger_mode"` +} + +type Result struct { + Messages []provider.Message `json:"messages"` + Metrics Metrics `json:"metrics"` + TranscriptID string `json:"transcript_id"` + TranscriptPath string `json:"transcript_path"` + Applied bool `json:"applied"` + ErrorMode ErrorMode `json:"error_mode"` +} + +type Runner interface { + Run(ctx context.Context, input Input) (Result, error) +} + +type Service struct { + now func() time.Time + randomToken func() (string, error) + userHomeDir func() (string, error) + mkdirAll func(path string, perm os.FileMode) error + writeFile func(name string, data []byte, perm os.FileMode) error + rename func(oldPath, newPath string) error + remove func(path string) error +} + +func NewRunner() *Service { + return &Service{ + now: time.Now, + randomToken: randomTranscriptToken, + userHomeDir: os.UserHomeDir, + mkdirAll: os.MkdirAll, + writeFile: os.WriteFile, + rename: os.Rename, + remove: os.Remove, + } +} + +func (s *Service) Run(ctx context.Context, input Input) (Result, error) { + if err := ctx.Err(); err != nil { + return Result{}, err + } + + cfg := normalizeCompactConfig(input.Config) + messages := cloneMessages(input.Messages) + + beforeChars := countMessageChars(messages) + base := Result{ + Messages: messages, + Applied: false, + ErrorMode: ErrorModeNone, + Metrics: Metrics{ + BeforeChars: beforeChars, + AfterChars: beforeChars, + SavedRatio: 0, + TriggerMode: string(input.Mode), + }, + } + + switch input.Mode { + case ModeMicro: + if !cfg.MicroEnabled || !hasMicroCandidate(messages, cfg) { + return base, nil + } + case ModeManual: + // manual compact is always evaluated when explicitly requested. + default: + return Result{}, fmt.Errorf("compact: unsupported mode %q", input.Mode) + } + + transcriptID, transcriptPath, err := s.saveTranscript(messages, strings.TrimSpace(input.SessionID), strings.TrimSpace(input.Workdir)) + if err != nil { + return Result{}, err + } + base.TranscriptID = transcriptID + base.TranscriptPath = transcriptPath + + var ( + next []provider.Message + applied bool + ) + + switch input.Mode { + case ModeMicro: + next, applied = microCompact(messages, cfg) + case ModeManual: + next, applied, err = manualCompact(messages, cfg) + if err != nil { + return Result{}, err + } + } + + afterChars := countMessageChars(next) + result := base + result.Messages = next + result.Applied = applied + result.Metrics.AfterChars = afterChars + if beforeChars > 0 { + result.Metrics.SavedRatio = float64(beforeChars-afterChars) / float64(beforeChars) + } + return result, nil +} + +func hasMicroCandidate(messages []provider.Message, cfg config.CompactConfig) bool { + toolIndices := make([]int, 0, len(messages)) + for i, message := range messages { + if message.Role == provider.RoleTool { + toolIndices = append(toolIndices, i) + } + } + if len(toolIndices) <= cfg.ToolResultKeepRecent { + return false + } + candidateCount := len(toolIndices) - cfg.ToolResultKeepRecent + for i := 0; i < candidateCount; i++ { + if len(messages[toolIndices[i]].Content) >= cfg.ToolResultPlaceholderMinChars { + return true + } + } + return false +} + +func microCompact(messages []provider.Message, cfg config.CompactConfig) ([]provider.Message, bool) { + next := cloneMessages(messages) + + toolIndices := make([]int, 0, len(next)) + for i, message := range next { + if message.Role == provider.RoleTool { + toolIndices = append(toolIndices, i) + } + } + + if len(toolIndices) <= cfg.ToolResultKeepRecent { + return next, false + } + + toolNameByCallID := buildToolNameIndex(next) + candidateCount := len(toolIndices) - cfg.ToolResultKeepRecent + applied := false + for i := 0; i < candidateCount; i++ { + idx := toolIndices[i] + message := next[idx] + if len(message.Content) < cfg.ToolResultPlaceholderMinChars { + continue + } + + toolName := strings.TrimSpace(toolNameByCallID[message.ToolCallID]) + if toolName == "" { + toolName = "unknown_tool" + } + placeholder := fmt.Sprintf("[Previous tool used: %s]", toolName) + if message.Content == placeholder { + continue + } + message.Content = placeholder + next[idx] = message + applied = true + } + + return next, applied +} + +func buildToolNameIndex(messages []provider.Message) map[string]string { + index := make(map[string]string) + for _, message := range messages { + if message.Role != provider.RoleAssistant || len(message.ToolCalls) == 0 { + continue + } + for _, call := range message.ToolCalls { + id := strings.TrimSpace(call.ID) + name := strings.TrimSpace(call.Name) + if id == "" || name == "" { + continue + } + index[id] = name + } + } + return index +} + +type span struct { + start int + end int +} + +func collectSpans(messages []provider.Message) []span { + spans := make([]span, 0, len(messages)) + for i := 0; i < len(messages); { + start := i + i++ + + if messages[start].Role == provider.RoleAssistant && len(messages[start].ToolCalls) > 0 { + for i < len(messages) && messages[i].Role == provider.RoleTool { + i++ + } + } + + spans = append(spans, span{start: start, end: i}) + } + return spans +} + +func manualCompact(messages []provider.Message, cfg config.CompactConfig) ([]provider.Message, bool, error) { + strategy := strings.ToLower(strings.TrimSpace(cfg.ManualStrategy)) + switch strategy { + case config.CompactManualStrategyKeepRecent: + return manualCompactKeepRecent(messages, cfg) + case config.CompactManualStrategyFullReplace: + return manualCompactFullReplace(messages, cfg) + default: + return nil, false, fmt.Errorf("compact: manual strategy %q is not supported", cfg.ManualStrategy) + } +} + +func manualCompactKeepRecent(messages []provider.Message, cfg config.CompactConfig) ([]provider.Message, bool, error) { + spans := collectSpans(messages) + if len(spans) <= cfg.ManualKeepRecentSpans { + return cloneMessages(messages), false, nil + } + + keepStart := spans[len(spans)-cfg.ManualKeepRecentSpans].start + removed := cloneMessages(messages[:keepStart]) + kept := cloneMessages(messages[keepStart:]) + + summary, err := buildSummary(removed, len(spans)-cfg.ManualKeepRecentSpans, cfg) + if err != nil { + return nil, false, err + } + + next := make([]provider.Message, 0, len(kept)+1) + next = append(next, provider.Message{Role: provider.RoleAssistant, Content: summary}) + next = append(next, kept...) + return next, true, nil +} + +func manualCompactFullReplace(messages []provider.Message, cfg config.CompactConfig) ([]provider.Message, bool, error) { + if len(messages) == 0 { + return nil, false, nil + } + spans := collectSpans(messages) + summary, err := buildSummary(cloneMessages(messages), len(spans), cfg) + if err != nil { + return nil, false, err + } + + return []provider.Message{{Role: provider.RoleAssistant, Content: summary}}, true, nil +} + +func buildSummary(removed []provider.Message, removedSpans int, cfg config.CompactConfig) (string, error) { + toolNames := make([]string, 0, 8) + seenTools := map[string]struct{}{} + for _, message := range removed { + for _, call := range message.ToolCalls { + name := strings.TrimSpace(call.Name) + if name == "" { + continue + } + if _, exists := seenTools[name]; exists { + continue + } + seenTools[name] = struct{}{} + toolNames = append(toolNames, name) + } + } + toolSummary := "none" + if len(toolNames) > 0 { + toolSummary = strings.Join(toolNames, ", ") + } + + summary := fmt.Sprintf(strings.Join([]string{ + "[compact_summary]", + "done:", + "- Archived %d historical spans (%d messages).", + "", + "in_progress:", + "- Continue from the retained recent context window.", + "", + "decisions:", + "- manual_strategy=%s", + "- manual_keep_recent_spans=%d", + "", + "code_changes:", + "- Older context outside the recent window was replaced by this summary.", + "- Historical tool calls in archived spans: %s", + "", + "constraints:", + "- Assistant tool_calls and tool_result pairs remain intact in retained spans.", + }, "\n"), removedSpans, len(removed), cfg.ManualStrategy, cfg.ManualKeepRecentSpans, toolSummary) + + return validateSummary(summary, cfg.MaxSummaryChars) +} + +func validateSummary(summary string, maxChars int) (string, error) { + summary = strings.TrimSpace(summary) + if maxChars > 0 && len(summary) > maxChars { + summary = strings.TrimSpace(summary[:maxChars]) + } + + hasDone := sectionHasContent(summary, "done") + hasInProgress := sectionHasContent(summary, "in_progress") + if !hasDone && !hasInProgress { + return "", errors.New("compact: summary requires done or in_progress content") + } + return summary, nil +} + +func sectionHasContent(summary string, section string) bool { + pattern := fmt.Sprintf(`(?ms)^%s:\s*\n\s*-\s+.+?(\n\w+?:|\z)`, regexp.QuoteMeta(section)) + re, err := regexp.Compile(pattern) + if err != nil { + return false + } + return re.MatchString(summary) +} + +type transcriptLine struct { + Index int `json:"index"` + Timestamp string `json:"timestamp"` + Role string `json:"role"` + Content string `json:"content"` + ToolCalls []provider.ToolCall `json:"tool_calls,omitempty"` + ToolCallID string `json:"tool_call_id,omitempty"` + IsError bool `json:"is_error,omitempty"` +} + +func (s *Service) saveTranscript(messages []provider.Message, sessionID string, workdir string) (string, string, error) { + home, err := s.userHomeDir() + if err != nil { + return "", "", fmt.Errorf("compact: resolve user home: %w", err) + } + + projectHash := hashProject(workdir) + dir := filepath.Join(home, ".neocode", "projects", projectHash, ".transcripts") + if err := s.mkdirAll(dir, 0o755); err != nil { + return "", "", fmt.Errorf("compact: create transcript dir: %w", err) + } + + sessionID = sanitizeID(sessionID) + if sessionID == "" { + sessionID = "draft" + } + tokenFn := s.randomToken + if tokenFn == nil { + tokenFn = randomTranscriptToken + } + randomToken, err := tokenFn() + if err != nil { + return "", "", fmt.Errorf("compact: generate transcript token: %w", err) + } + transcriptID := fmt.Sprintf("transcript_%d_%s_%s", s.now().UnixNano(), randomToken, sessionID) + transcriptPath := filepath.Join(dir, transcriptID+".jsonl") + tmpPath := transcriptPath + ".tmp" + + now := s.now().UTC().Format(time.RFC3339Nano) + var builder strings.Builder + for i, message := range messages { + line := transcriptLine{ + Index: i, + Timestamp: now, + Role: message.Role, + Content: message.Content, + ToolCalls: append([]provider.ToolCall(nil), message.ToolCalls...), + ToolCallID: message.ToolCallID, + IsError: message.IsError, + } + payload, err := json.Marshal(line) + if err != nil { + return "", "", fmt.Errorf("compact: marshal transcript line: %w", err) + } + builder.Write(payload) + builder.WriteByte('\n') + } + + if err := s.writeFile(tmpPath, []byte(builder.String()), transcriptFileMode()); err != nil { + return "", "", fmt.Errorf("compact: write transcript: %w", err) + } + if err := s.rename(tmpPath, transcriptPath); err != nil { + _ = s.remove(tmpPath) + return "", "", fmt.Errorf("compact: commit transcript: %w", err) + } + + return transcriptID, transcriptPath, nil +} + +func transcriptFileMode() os.FileMode { + if goruntime.GOOS == "windows" { + return 0o644 + } + return 0o600 +} + +func randomTranscriptToken() (string, error) { + entropy := make([]byte, 4) + if _, err := cryptorand.Read(entropy); err != nil { + return "", err + } + return hex.EncodeToString(entropy), nil +} + +func hashProject(workdir string) string { + clean := strings.TrimSpace(filepath.Clean(workdir)) + if clean == "" { + clean = "unknown" + } + sum := sha1.Sum([]byte(strings.ToLower(clean))) + return hex.EncodeToString(sum[:8]) +} + +var nonIDChars = regexp.MustCompile(`[^a-zA-Z0-9_-]+`) + +func sanitizeID(value string) string { + value = strings.TrimSpace(value) + if value == "" { + return "" + } + return nonIDChars.ReplaceAllString(value, "_") +} + +func cloneMessages(messages []provider.Message) []provider.Message { + if len(messages) == 0 { + return nil + } + out := make([]provider.Message, 0, len(messages)) + for _, message := range messages { + next := message + next.ToolCalls = append([]provider.ToolCall(nil), message.ToolCalls...) + out = append(out, next) + } + return out +} + +func countMessageChars(messages []provider.Message) int { + total := 0 + for _, message := range messages { + total += len(message.Role) + total += len(message.Content) + total += len(message.ToolCallID) + for _, call := range message.ToolCalls { + total += len(call.ID) + total += len(call.Name) + total += len(call.Arguments) + } + } + return total +} + +func normalizeCompactConfig(cfg config.CompactConfig) config.CompactConfig { + defaults := config.Default().Context.Compact + cfg.ApplyDefaults(defaults) + if strings.TrimSpace(cfg.ManualStrategy) == "" { + cfg.ManualStrategy = config.CompactManualStrategyKeepRecent + } + return cfg +} diff --git a/internal/context/compact/runner_test.go b/internal/context/compact/runner_test.go new file mode 100644 index 00000000..8b1608f5 --- /dev/null +++ b/internal/context/compact/runner_test.go @@ -0,0 +1,337 @@ +package compact + +import ( + "context" + "errors" + "os" + "path/filepath" + goruntime "runtime" + "strings" + "testing" + "time" + + "neo-code/internal/config" + "neo-code/internal/provider" +) + +func TestMicroCompactReplacesOnlyOldLongToolResults(t *testing.T) { + t.Parallel() + + runner := NewRunner() + home := t.TempDir() + runner.userHomeDir = func() (string, error) { return home, nil } + + messages := []provider.Message{ + {Role: provider.RoleUser, Content: "start"}, + {Role: provider.RoleAssistant, ToolCalls: []provider.ToolCall{{ID: "call-1", Name: "filesystem_read_file", Arguments: "{}"}}}, + {Role: provider.RoleTool, ToolCallID: "call-1", Content: strings.Repeat("A", 40)}, + {Role: provider.RoleAssistant, ToolCalls: []provider.ToolCall{{ID: "call-2", Name: "filesystem_grep", Arguments: "{}"}}}, + {Role: provider.RoleTool, ToolCallID: "call-2", Content: "short"}, + {Role: provider.RoleAssistant, ToolCalls: []provider.ToolCall{{ID: "call-3", Name: "bash", Arguments: "{}"}}}, + {Role: provider.RoleTool, ToolCallID: "call-3", Content: strings.Repeat("B", 50)}, + } + + result, err := runner.Run(context.Background(), Input{ + Mode: ModeMicro, + SessionID: "session-a", + Workdir: t.TempDir(), + Messages: messages, + Config: config.CompactConfig{ + MicroEnabled: true, + ToolResultKeepRecent: 1, + ToolResultPlaceholderMinChars: 10, + ManualStrategy: config.CompactManualStrategyKeepRecent, + ManualKeepRecentSpans: 6, + MaxSummaryChars: 1200, + }, + }) + if err != nil { + t.Fatalf("Run() error = %v", err) + } + if !result.Applied { + t.Fatalf("expected micro compact applied") + } + if !strings.Contains(result.Messages[2].Content, "filesystem_read_file") { + t.Fatalf("expected first old tool result to be replaced, got %q", result.Messages[2].Content) + } + if result.Messages[4].Content != "short" { + t.Fatalf("expected short old tool result unchanged, got %q", result.Messages[4].Content) + } + if result.Messages[6].Content == "[Previous tool used: bash]" { + t.Fatalf("expected recent tool result to be retained") + } + if result.TranscriptID == "" || result.TranscriptPath == "" { + t.Fatalf("expected transcript metadata, got %+v", result) + } + if _, err := os.Stat(result.TranscriptPath); err != nil { + t.Fatalf("expected transcript file: %v", err) + } +} + +func TestMicroCompactFallsBackToUnknownTool(t *testing.T) { + t.Parallel() + + runner := NewRunner() + home := t.TempDir() + runner.userHomeDir = func() (string, error) { return home, nil } + + messages := []provider.Message{ + {Role: provider.RoleUser, Content: "start"}, + {Role: provider.RoleTool, ToolCallID: "missing-call", Content: strings.Repeat("X", 24)}, + {Role: provider.RoleTool, ToolCallID: "recent-call", Content: strings.Repeat("Y", 24)}, + } + + result, err := runner.Run(context.Background(), Input{ + Mode: ModeMicro, + SessionID: "session-b", + Workdir: t.TempDir(), + Messages: messages, + Config: config.CompactConfig{ + MicroEnabled: true, + ToolResultKeepRecent: 1, + ToolResultPlaceholderMinChars: 10, + ManualStrategy: config.CompactManualStrategyKeepRecent, + ManualKeepRecentSpans: 6, + MaxSummaryChars: 1200, + }, + }) + if err != nil { + t.Fatalf("Run() error = %v", err) + } + if got := result.Messages[1].Content; got != "[Previous tool used: unknown_tool]" { + t.Fatalf("expected unknown tool placeholder, got %q", got) + } +} + +func TestManualCompactAddsSummaryAndKeepsRecentSpans(t *testing.T) { + t.Parallel() + + runner := NewRunner() + home := t.TempDir() + runner.userHomeDir = func() (string, error) { return home, nil } + + messages := []provider.Message{ + {Role: provider.RoleUser, Content: "old requirement"}, + {Role: provider.RoleAssistant, ToolCalls: []provider.ToolCall{{ID: "call-old", Name: "filesystem_grep", Arguments: "{}"}}}, + {Role: provider.RoleTool, ToolCallID: "call-old", Content: "old result"}, + {Role: provider.RoleAssistant, Content: "latest answer"}, + } + + result, err := runner.Run(context.Background(), Input{ + Mode: ModeManual, + SessionID: "session-c", + Workdir: t.TempDir(), + Messages: messages, + Config: config.CompactConfig{ + MicroEnabled: false, + ToolResultKeepRecent: 3, + ToolResultPlaceholderMinChars: 100, + ManualStrategy: config.CompactManualStrategyKeepRecent, + ManualKeepRecentSpans: 1, + MaxSummaryChars: 1200, + }, + }) + if err != nil { + t.Fatalf("Run() error = %v", err) + } + if !result.Applied { + t.Fatalf("expected manual compact applied") + } + if len(result.Messages) != 2 { + t.Fatalf("expected summary + 1 kept span, got %d", len(result.Messages)) + } + summary := result.Messages[0] + if summary.Role != provider.RoleAssistant { + t.Fatalf("expected summary role assistant, got %q", summary.Role) + } + for _, section := range []string{"done:", "in_progress:", "decisions:", "code_changes:", "constraints:"} { + if !strings.Contains(summary.Content, section) { + t.Fatalf("expected summary to include section %q, got %q", section, summary.Content) + } + } + if result.Messages[1].Content != "latest answer" { + t.Fatalf("expected newest span kept, got %+v", result.Messages[1]) + } +} + +func TestManualCompactWritesTranscriptJSONL(t *testing.T) { + t.Parallel() + + runner := NewRunner() + home := t.TempDir() + runner.userHomeDir = func() (string, error) { return home, nil } + + result, err := runner.Run(context.Background(), Input{ + Mode: ModeManual, + SessionID: "session-jsonl", + Workdir: filepath.Join(home, "workspace"), + Messages: []provider.Message{ + {Role: provider.RoleUser, Content: "hello"}, + }, + Config: config.CompactConfig{ + ManualStrategy: config.CompactManualStrategyKeepRecent, + ManualKeepRecentSpans: 6, + ToolResultKeepRecent: 3, + ToolResultPlaceholderMinChars: 100, + MaxSummaryChars: 1200, + }, + }) + if err != nil { + t.Fatalf("Run() error = %v", err) + } + data, err := os.ReadFile(result.TranscriptPath) + if err != nil { + t.Fatalf("read transcript: %v", err) + } + if !strings.Contains(string(data), `"role":"user"`) { + t.Fatalf("expected jsonl content, got %q", string(data)) + } + if !strings.Contains(filepath.ToSlash(result.TranscriptPath), "/.neocode/projects/") { + t.Fatalf("expected transcript path under .neocode/projects, got %q", result.TranscriptPath) + } + if !strings.HasPrefix(result.TranscriptID, "transcript_") { + t.Fatalf("unexpected transcript id: %q", result.TranscriptID) + } + if goruntime.GOOS != "windows" { + info, err := os.Stat(result.TranscriptPath) + if err != nil { + t.Fatalf("stat transcript: %v", err) + } + if got := info.Mode().Perm(); got != 0o600 { + t.Fatalf("expected transcript mode 0600, got %04o", got) + } + } +} + +func TestManualCompactFailsWhenTranscriptWriteFails(t *testing.T) { + t.Parallel() + + runner := NewRunner() + runner.userHomeDir = func() (string, error) { return t.TempDir(), nil } + runner.mkdirAll = func(path string, perm os.FileMode) error { + return errors.New("disk full") + } + + _, err := runner.Run(context.Background(), Input{ + Mode: ModeManual, + SessionID: "session-fail", + Workdir: t.TempDir(), + Messages: []provider.Message{ + {Role: provider.RoleUser, Content: "hello"}, + }, + Config: config.CompactConfig{ + ManualStrategy: config.CompactManualStrategyKeepRecent, + ManualKeepRecentSpans: 6, + ToolResultKeepRecent: 3, + ToolResultPlaceholderMinChars: 100, + MaxSummaryChars: 1200, + }, + }) + if err == nil || !strings.Contains(err.Error(), "disk full") { + t.Fatalf("expected transcript write failure, got %v", err) + } +} + +func TestManualCompactFullReplaceRewritesAllMessages(t *testing.T) { + t.Parallel() + + runner := NewRunner() + home := t.TempDir() + runner.userHomeDir = func() (string, error) { return home, nil } + + messages := []provider.Message{ + {Role: provider.RoleUser, Content: "old requirement"}, + {Role: provider.RoleAssistant, ToolCalls: []provider.ToolCall{{ID: "call-old", Name: "filesystem_grep", Arguments: "{}"}}}, + {Role: provider.RoleTool, ToolCallID: "call-old", Content: "old result"}, + {Role: provider.RoleAssistant, Content: "latest answer"}, + } + + result, err := runner.Run(context.Background(), Input{ + Mode: ModeManual, + SessionID: "session-full-replace", + Workdir: t.TempDir(), + Messages: messages, + Config: config.CompactConfig{ + MicroEnabled: true, + ToolResultKeepRecent: 3, + ToolResultPlaceholderMinChars: 100, + ManualStrategy: config.CompactManualStrategyFullReplace, + ManualKeepRecentSpans: 6, + MaxSummaryChars: 1200, + }, + }) + if err != nil { + t.Fatalf("Run() error = %v", err) + } + if !result.Applied { + t.Fatalf("expected full_replace compact applied") + } + if len(result.Messages) != 1 { + t.Fatalf("expected single summary message, got %d", len(result.Messages)) + } + if result.Messages[0].Role != provider.RoleAssistant { + t.Fatalf("expected summary role assistant, got %q", result.Messages[0].Role) + } + for _, section := range []string{"done:", "in_progress:", "decisions:", "code_changes:", "constraints:"} { + if !strings.Contains(result.Messages[0].Content, section) { + t.Fatalf("expected summary section %q, got %q", section, result.Messages[0].Content) + } + } +} + +func TestSaveTranscriptUsesUniqueIDWithinSameTimestamp(t *testing.T) { + t.Parallel() + + runner := NewRunner() + home := t.TempDir() + runner.userHomeDir = func() (string, error) { return home, nil } + fixedNow := time.Unix(1712052000, 123456789) + runner.now = func() time.Time { return fixedNow } + tokenSeq := []string{"a1b2c3d4", "b2c3d4e5"} + runner.randomToken = func() (string, error) { + if len(tokenSeq) == 0 { + return "", errors.New("empty token sequence") + } + next := tokenSeq[0] + tokenSeq = tokenSeq[1:] + return next, nil + } + + input := Input{ + Mode: ModeManual, + SessionID: "session-dup-safe", + Workdir: t.TempDir(), + Messages: []provider.Message{ + {Role: provider.RoleUser, Content: "hello"}, + {Role: provider.RoleAssistant, Content: "world"}, + }, + Config: config.CompactConfig{ + ManualStrategy: config.CompactManualStrategyFullReplace, + ManualKeepRecentSpans: 6, + ToolResultKeepRecent: 3, + ToolResultPlaceholderMinChars: 100, + MaxSummaryChars: 1200, + }, + } + + first, err := runner.Run(context.Background(), input) + if err != nil { + t.Fatalf("first Run() error = %v", err) + } + second, err := runner.Run(context.Background(), input) + if err != nil { + t.Fatalf("second Run() error = %v", err) + } + if first.TranscriptID == second.TranscriptID { + t.Fatalf("expected distinct transcript ids, got %q", first.TranscriptID) + } + if first.TranscriptPath == second.TranscriptPath { + t.Fatalf("expected distinct transcript paths, got %q", first.TranscriptPath) + } + if _, err := os.Stat(first.TranscriptPath); err != nil { + t.Fatalf("first transcript file missing: %v", err) + } + if _, err := os.Stat(second.TranscriptPath); err != nil { + t.Fatalf("second transcript file missing: %v", err) + } +} diff --git a/internal/runtime/events.go b/internal/runtime/events.go index 0e2ce006..27cd73b7 100644 --- a/internal/runtime/events.go +++ b/internal/runtime/events.go @@ -36,4 +36,12 @@ const ( // EventProviderRetry is emitted when runtime retries a provider call due to // a retryable error (e.g. 429, 5xx). Payload is a human-readable message. EventProviderRetry EventType = "provider_retry" + // EventCompactStart is emitted when a compact cycle starts. + EventCompactStart EventType = "compact_start" + // EventCompactDone is emitted when a compact cycle completes. + EventCompactDone EventType = "compact_done" + // EventCompactError is emitted when compact fails. + EventCompactError EventType = "compact_error" + // EventMicroCompactApplied is emitted when micro compact has rewritten context. + EventMicroCompactApplied EventType = "micro_compact_applied" ) diff --git a/internal/runtime/runtime.go b/internal/runtime/runtime.go index a9ea5dac..1fe8cef1 100644 --- a/internal/runtime/runtime.go +++ b/internal/runtime/runtime.go @@ -12,6 +12,7 @@ import ( "neo-code/internal/config" agentcontext "neo-code/internal/context" + contextcompact "neo-code/internal/context/compact" "neo-code/internal/provider" "neo-code/internal/tools" ) @@ -30,6 +31,7 @@ const ( type Runtime interface { Run(ctx context.Context, input UserInput) error + Compact(ctx context.Context, input CompactInput) (CompactResult, error) CancelActiveRun() bool Events() <-chan RuntimeEvent ListSessions(ctx context.Context) ([]SessionSummary, error) @@ -42,6 +44,36 @@ type UserInput struct { Content string } +type CompactInput struct { + SessionID string + RunID string +} + +type CompactResult struct { + Applied bool + BeforeChars int + AfterChars int + SavedRatio float64 + TriggerMode string + TranscriptID string + TranscriptPath string +} + +type CompactDonePayload struct { + Applied bool `json:"applied"` + BeforeChars int `json:"before_chars"` + AfterChars int `json:"after_chars"` + SavedRatio float64 `json:"saved_ratio"` + TriggerMode string `json:"trigger_mode"` + TranscriptID string `json:"transcript_id"` + TranscriptPath string `json:"transcript_path"` +} + +type CompactErrorPayload struct { + TriggerMode string `json:"trigger_mode"` + Message string `json:"message"` +} + type ProviderFactory interface { Build(ctx context.Context, cfg config.ResolvedProviderConfig) (provider.Provider, error) } @@ -52,11 +84,12 @@ type Service struct { toolManager tools.Manager // 工具管理器，统一工具 schema 暴露与执行入口。 providerFactory ProviderFactory // Provider 工厂接口，根据配置动态创建具体的 provider 实例。 contextBuilder agentcontext.Builder // 上下文构建器，负责组装 system prompt 与本轮发给模型的消息上下文。 - events chan RuntimeEvent // 事件通道，Runtime 在运行过程中产生的事件都通过该通道发送给 TUI 层消费和展示。 - runMu sync.Mutex // 运行互斥锁，保证同一时间只有一个 Run 在执行。 - activeRunToken uint64 // 当前活跃运行的令牌标识，用于标记正在执行的 Run 实例。 - nextRunToken uint64 // 下一个运行令牌的递增计数器，用于区分不同 Run 的生命周期。 - activeRunCancel context.CancelFunc // 当前活跃 Run 的取消函数。 + compactRunner contextcompact.Runner + events chan RuntimeEvent // 事件通道，Runtime 在运行过程中产生的事件都通过该通道发送给 TUI 层消费和展示。 + runMu sync.Mutex // 运行互斥锁，保证同一时间只有一个 Run 在执行。 + activeRunToken uint64 // 当前活跃运行的令牌标识，用于标记正在执行的 Run 实例。 + nextRunToken uint64 // 下一个运行令牌的递增计数器，用于区分不同 Run 的生命周期。 + activeRunCancel context.CancelFunc // 当前活跃 Run 的取消函数。 } func NewWithFactory( @@ -82,6 +115,7 @@ func NewWithFactory( toolManager: toolManager, providerFactory: providerFactory, contextBuilder: contextBuilder, + compactRunner: contextcompact.NewRunner(), events: make(chan RuntimeEvent, 128), } } @@ -131,6 +165,11 @@ func (s *Service) Run(ctx context.Context, input UserInput) error { return err } + session, _, err = s.runCompactForSession(ctx, input.RunID, session, cfg, contextcompact.ModeMicro, false) + if err != nil { + return s.handleRunError(ctx, input.RunID, session.ID, err) + } + builtContext, err := s.contextBuilder.Build(ctx, agentcontext.BuildInput{ Messages: session.Messages, Metadata: agentcontext.Metadata{ @@ -246,6 +285,36 @@ func (s *Service) Run(ctx context.Context, input UserInput) error { } } +func (s *Service) Compact(ctx context.Context, input CompactInput) (CompactResult, error) { + if err := ctx.Err(); err != nil { + return CompactResult{}, err + } + if strings.TrimSpace(input.SessionID) == "" { + return CompactResult{}, errors.New("runtime: compact session_id is empty") + } + + cfg := s.configManager.Get() + session, err := s.sessionStore.Load(ctx, input.SessionID) + if err != nil { + return CompactResult{}, err + } + + session, result, err := s.runCompactForSession(ctx, input.RunID, session, cfg, contextcompact.ModeManual, true) + if err != nil { + return CompactResult{}, err + } + + return CompactResult{ + Applied: result.Applied, + BeforeChars: result.Metrics.BeforeChars, + AfterChars: result.Metrics.AfterChars, + SavedRatio: result.Metrics.SavedRatio, + TriggerMode: result.Metrics.TriggerMode, + TranscriptID: result.TranscriptID, + TranscriptPath: result.TranscriptPath, + }, nil +} + func (s *Service) CancelActiveRun() bool { s.runMu.Lock() cancel := s.activeRunCancel @@ -364,6 +433,72 @@ func (s *Service) handleRunError(ctx context.Context, runID string, sessionID st return err } +func (s *Service) runCompactForSession( + ctx context.Context, + runID string, + session Session, + cfg config.Config, + mode contextcompact.Mode, + failOnError bool, +) (Session, contextcompact.Result, error) { + if s.compactRunner == nil { + return session, contextcompact.Result{}, nil + } + + originalMessages := append([]provider.Message(nil), session.Messages...) + s.emit(ctx, EventCompactStart, runID, session.ID, string(mode)) + + result, err := s.compactRunner.Run(ctx, contextcompact.Input{ + Mode: mode, + SessionID: session.ID, + Workdir: cfg.Workdir, + Messages: session.Messages, + Config: cfg.Context.Compact, + }) + if err != nil { + s.emit(ctx, EventCompactError, runID, session.ID, CompactErrorPayload{ + TriggerMode: string(mode), + Message: err.Error(), + }) + if failOnError { + return session, contextcompact.Result{}, err + } + return session, contextcompact.Result{}, nil + } + + if result.Applied { + session.Messages = append([]provider.Message(nil), result.Messages...) + session.UpdatedAt = time.Now() + if err := s.sessionStore.Save(ctx, &session); err != nil { + s.emit(ctx, EventCompactError, runID, session.ID, CompactErrorPayload{ + TriggerMode: string(mode), + Message: err.Error(), + }) + session.Messages = originalMessages + if failOnError { + return session, contextcompact.Result{}, err + } + return session, contextcompact.Result{}, nil + } + } + + donePayload := CompactDonePayload{ + Applied: result.Applied, + BeforeChars: result.Metrics.BeforeChars, + AfterChars: result.Metrics.AfterChars, + SavedRatio: result.Metrics.SavedRatio, + TriggerMode: result.Metrics.TriggerMode, + TranscriptID: result.TranscriptID, + TranscriptPath: result.TranscriptPath, + } + s.emit(ctx, EventCompactDone, runID, session.ID, donePayload) + if mode == contextcompact.ModeMicro && result.Applied { + s.emit(ctx, EventMicroCompactApplied, runID, session.ID, donePayload) + } + + return session, result, nil +} + // isRetryableProviderError 判断 error 是否为可重试的 ProviderError。 func isRetryableProviderError(err error) bool { var pErr *provider.ProviderError diff --git a/internal/runtime/runtime_test.go b/internal/runtime/runtime_test.go index f75a4cf8..9ceba70e 100644 --- a/internal/runtime/runtime_test.go +++ b/internal/runtime/runtime_test.go @@ -11,6 +11,7 @@ import ( "neo-code/internal/config" agentcontext "neo-code/internal/context" + contextcompact "neo-code/internal/context/compact" "neo-code/internal/provider" "neo-code/internal/provider/builtin" "neo-code/internal/tools" @@ -203,6 +204,23 @@ type stubToolManager struct { lastInput tools.ToolCallInput } +type stubCompactRunner struct { + runFn func(ctx context.Context, input contextcompact.Input) (contextcompact.Result, error) + calls []contextcompact.Input + result contextcompact.Result + err error +} + +func (r *stubCompactRunner) Run(ctx context.Context, input contextcompact.Input) (contextcompact.Result, error) { + cloned := input + cloned.Messages = append([]provider.Message(nil), input.Messages...) + r.calls = append(r.calls, cloned) + if r.runFn != nil { + return r.runFn(ctx, input) + } + return r.result, r.err +} + func (m *stubToolManager) ListAvailableSpecs(ctx context.Context, input tools.SpecListInput) ([]provider.ToolSpec, error) { m.listCalls++ if err := ctx.Err(); err != nil { @@ -1231,6 +1249,204 @@ func TestServiceConstructorsAndDelegates(t *testing.T) { } } +func TestServiceRunAppliesMicroCompactBeforeContextBuild(t *testing.T) { + manager := newRuntimeConfigManager(t) + if err := manager.Update(context.Background(), func(cfg *config.Config) error { + cfg.Context.Compact.MicroEnabled = true + return nil + }); err != nil { + t.Fatalf("enable micro compact: %v", err) + } + + store := newMemoryStore() + session := newSession("existing") + session.ID = "session-micro" + session.Messages = []provider.Message{ + {Role: provider.RoleUser, Content: "history"}, + } + store.sessions[session.ID] = cloneSession(session) + + registry := tools.NewRegistry() + registry.Register(&stubTool{name: "filesystem_read_file", content: "ok"}) + + builder := &stubContextBuilder{} + scripted := &scriptedProvider{ + responses: []provider.ChatResponse{ + { + Message: provider.Message{Role: provider.RoleAssistant, Content: "done"}, + }, + }, + } + service := NewWithFactory(manager, registry, store, &scriptedProviderFactory{provider: scripted}, builder) + compactRunner := &stubCompactRunner{ + runFn: func(ctx context.Context, input contextcompact.Input) (contextcompact.Result, error) { + next := append([]provider.Message(nil), input.Messages...) + next = append(next, provider.Message{Role: provider.RoleAssistant, Content: "micro-summary"}) + return contextcompact.Result{ + Messages: next, + Applied: true, + Metrics: contextcompact.Metrics{ + BeforeChars: 10, + AfterChars: 6, + SavedRatio: 0.4, + TriggerMode: string(contextcompact.ModeMicro), + }, + TranscriptID: "transcript_1_session-micro", + TranscriptPath: "/tmp/transcript.jsonl", + }, nil + }, + } + service.compactRunner = compactRunner + + if err := service.Run(context.Background(), UserInput{ + SessionID: session.ID, + RunID: "run-micro", + Content: "latest", + }); err != nil { + t.Fatalf("Run() error = %v", err) + } + + if len(compactRunner.calls) == 0 || compactRunner.calls[0].Mode != contextcompact.ModeMicro { + t.Fatalf("expected micro compact call, got %+v", compactRunner.calls) + } + if len(builder.lastInput.Messages) == 0 || builder.lastInput.Messages[len(builder.lastInput.Messages)-1].Content != "micro-summary" { + t.Fatalf("expected compacted messages to be passed into context builder, got %+v", builder.lastInput.Messages) + } + + events := collectRuntimeEvents(service.Events()) + assertEventSequence(t, events, []EventType{EventCompactStart, EventCompactDone, EventMicroCompactApplied, EventAgentDone}) + donePayload, ok := findEventPayload[CompactDonePayload](events, EventCompactDone) + if !ok || !donePayload.Applied || donePayload.TriggerMode != string(contextcompact.ModeMicro) { + t.Fatalf("expected compact_done payload with applied micro mode, got %+v", donePayload) + } +} + +func TestServiceRunMicroCompactFailureDoesNotBlockMainLoop(t *testing.T) { + manager := newRuntimeConfigManager(t) + if err := manager.Update(context.Background(), func(cfg *config.Config) error { + cfg.Context.Compact.MicroEnabled = true + return nil + }); err != nil { + t.Fatalf("enable micro compact: %v", err) + } + + store := newMemoryStore() + registry := tools.NewRegistry() + registry.Register(&stubTool{name: "filesystem_read_file", content: "ok"}) + + scripted := &scriptedProvider{ + responses: []provider.ChatResponse{ + { + Message: provider.Message{Role: provider.RoleAssistant, Content: "still works"}, + }, + }, + } + service := NewWithFactory(manager, registry, store, &scriptedProviderFactory{provider: scripted}, nil) + service.compactRunner = &stubCompactRunner{err: errors.New("compact failed")} + + if err := service.Run(context.Background(), UserInput{RunID: "run-micro-fail", Content: "hello"}); err != nil { + t.Fatalf("expected run to continue when micro compact fails, got %v", err) + } + + events := collectRuntimeEvents(service.Events()) + assertEventSequence(t, events, []EventType{EventCompactStart, EventCompactError, EventAgentDone}) + assertNoEventType(t, events, EventError) +} + +func TestServiceCompactManualAppliesAndPersists(t *testing.T) { + manager := newRuntimeConfigManager(t) + store := newMemoryStore() + session := newSession("manual") + session.ID = "session-manual" + session.Messages = []provider.Message{ + {Role: provider.RoleUser, Content: "before"}, + } + store.sessions[session.ID] = cloneSession(session) + + registry := tools.NewRegistry() + registry.Register(&stubTool{name: "filesystem_read_file", content: "ok"}) + + service := NewWithFactory(manager, registry, store, &scriptedProviderFactory{provider: &scriptedProvider{}}, nil) + service.compactRunner = &stubCompactRunner{ + result: contextcompact.Result{ + Messages: []provider.Message{ + {Role: provider.RoleAssistant, Content: "[compact_summary]\\ndone:\\n- ok\\n\\nin_progress:\\n- continue"}, + {Role: provider.RoleAssistant, Content: "latest"}, + }, + Applied: true, + Metrics: contextcompact.Metrics{ + BeforeChars: 80, + AfterChars: 30, + SavedRatio: 0.625, + TriggerMode: string(contextcompact.ModeManual), + }, + TranscriptID: "transcript_manual", + TranscriptPath: "/tmp/manual.jsonl", + }, + } + + result, err := service.Compact(context.Background(), CompactInput{ + SessionID: session.ID, + RunID: "run-manual", + }) + if err != nil { + t.Fatalf("Compact() error = %v", err) + } + if !result.Applied || result.BeforeChars != 80 || result.AfterChars != 30 { + t.Fatalf("unexpected compact result: %+v", result) + } + + saved, err := store.Load(context.Background(), session.ID) + if err != nil { + t.Fatalf("load compacted session: %v", err) + } + if len(saved.Messages) != 2 || !strings.Contains(saved.Messages[0].Content, "compact_summary") { + t.Fatalf("expected persisted compacted messages, got %+v", saved.Messages) + } + + events := collectRuntimeEvents(service.Events()) + assertEventSequence(t, events, []EventType{EventCompactStart, EventCompactDone}) + donePayload, ok := findEventPayload[CompactDonePayload](events, EventCompactDone) + if !ok || !donePayload.Applied || donePayload.TriggerMode != string(contextcompact.ModeManual) { + t.Fatalf("expected compact_done payload with applied manual mode, got %+v", donePayload) + } +} + +func TestServiceCompactManualFailureReturnsError(t *testing.T) { + manager := newRuntimeConfigManager(t) + store := newMemoryStore() + session := newSession("manual-fail") + session.ID = "session-manual-fail" + session.Messages = []provider.Message{{Role: provider.RoleUser, Content: "before"}} + store.sessions[session.ID] = cloneSession(session) + + registry := tools.NewRegistry() + registry.Register(&stubTool{name: "filesystem_read_file", content: "ok"}) + + service := NewWithFactory(manager, registry, store, &scriptedProviderFactory{provider: &scriptedProvider{}}, nil) + service.compactRunner = &stubCompactRunner{err: errors.New("manual compact failed")} + + _, err := service.Compact(context.Background(), CompactInput{ + SessionID: session.ID, + RunID: "run-manual-fail", + }) + if err == nil || !strings.Contains(err.Error(), "manual compact failed") { + t.Fatalf("expected compact failure, got %v", err) + } + + saved, err := store.Load(context.Background(), session.ID) + if err != nil { + t.Fatalf("load original session: %v", err) + } + if len(saved.Messages) != 1 || saved.Messages[0].Content != "before" { + t.Fatalf("expected original session untouched, got %+v", saved.Messages) + } + + events := collectRuntimeEvents(service.Events()) + assertEventSequence(t, events, []EventType{EventCompactStart, EventCompactError}) + assertNoEventType(t, events, EventCompactDone) +} + func newRuntimeConfigManager(t *testing.T) *config.Manager { t.Helper() restoreRuntimeEnv(t, config.OpenAIDefaultAPIKeyEnv) @@ -1312,6 +1528,21 @@ func assertNoEventType(t *testing.T, events []RuntimeEvent, unexpected EventType } } +func findEventPayload[T any](events []RuntimeEvent, kind EventType) (T, bool) { + var zero T + for _, event := range events { + if event.Type != kind { + continue + } + payload, ok := event.Payload.(T) + if !ok { + return zero, false + } + return payload, true + } + return zero, false +} + func assertEventsRunID(t *testing.T, events []RuntimeEvent, runID string) { t.Helper() for _, event := range events { diff --git a/internal/tui/commands.go b/internal/tui/commands.go index fa16e18a..7d03360d 100644 --- a/internal/tui/commands.go +++ b/internal/tui/commands.go @@ -17,6 +17,7 @@ const ( slashCommandHelp = "/help" slashCommandExit = "/exit" slashCommandClear = "/clear" + slashCommandCompact = "/compact" slashCommandStatus = "/status" slashCommandProvider = "/provider" slashCommandModelPick = "/model" @@ -24,6 +25,7 @@ const ( slashUsageHelp = "/help" slashUsageExit = "/exit" slashUsageClear = "/clear" + slashUsageCompact = "/compact" slashUsageStatus = "/status" slashUsageProvider = "/provider" slashUsageModel = "/model" @@ -58,6 +60,7 @@ const ( statusApplyingCommand = "Applying local command" statusRunningCommand = "Running command" statusCommandDone = "Command finished" + statusCompacting = "Compacting context" statusChooseProvider = "Choose a provider" statusChooseModel = "Choose a model" @@ -108,6 +111,7 @@ type statusSnapshot struct { var builtinSlashCommands = []slashCommand{ {Usage: slashUsageHelp, Description: "Show slash command help"}, {Usage: slashUsageClear, Description: "Clear the current draft transcript"}, + {Usage: slashUsageCompact, Description: "Compact the current session context"}, {Usage: slashUsageStatus, Description: "Show current session and agent status"}, {Usage: slashUsageProvider, Description: "Open the interactive provider picker"}, {Usage: slashUsageModel, Description: "Open the interactive model picker"}, diff --git a/internal/tui/commands_test.go b/internal/tui/commands_test.go index 903bdaa9..32c0747d 100644 --- a/internal/tui/commands_test.go +++ b/internal/tui/commands_test.go @@ -25,6 +25,7 @@ func TestExecuteLocalCommand(t *testing.T) { for _, want := range []string{ slashUsageHelp, slashUsageClear, + slashUsageCompact, slashUsageStatus, slashUsageProvider, slashUsageModel, diff --git a/internal/tui/update.go b/internal/tui/update.go index 5d8d82a9..8af49eb8 100644 --- a/internal/tui/update.go +++ b/internal/tui/update.go @@ -27,6 +27,9 @@ type modelCatalogRefreshMsg struct { models []provider.ModelDescriptor err error } +type compactFinishedMsg struct { + err error +} type localCommandResultMsg struct { notice string err error @@ -103,6 +106,26 @@ func (a App) Update(msg tea.Msg) (tea.Model, tea.Cmd) { a.syncConfigState(cfg) a.selectCurrentModel(cfg.CurrentModel) return a, tea.Batch(cmds...) + case compactFinishedMsg: + // compact feedback should primarily come from runtime compact_done/error events. + if typed.err != nil && strings.TrimSpace(a.state.ExecutionError) == "" { + a.state.ExecutionError = typed.err.Error() + a.state.StatusText = typed.err.Error() + } + if err := a.refreshSessions(); err != nil { + a.state.ExecutionError = err.Error() + a.state.StatusText = err.Error() + a.appendInlineMessage(roleError, err.Error()) + } + if err := a.refreshMessages(); err != nil && strings.TrimSpace(a.state.ActiveSessionID) != "" { + a.state.ExecutionError = err.Error() + a.state.StatusText = err.Error() + a.appendInlineMessage(roleError, err.Error()) + } + a.syncActiveSessionTitle() + a.rebuildTranscript() + a.transcript.GotoBottom() + return a, tea.Batch(cmds...) case localCommandResultMsg: if typed.err != nil { a.state.ExecutionError = typed.err.Error() @@ -525,6 +548,50 @@ func (a *App) handleRuntimeEvent(event agentruntime.RuntimeEvent) bool { a.state.StatusText = statusThinking a.appendActivity("provider", "Retrying provider call", payload, false) } + case agentruntime.EventCompactDone: + payload, ok := event.Payload.(agentruntime.CompactDonePayload) + if !ok { + return transcriptDirty + } + // Skip passive micro checks that did not rewrite context to avoid transcript noise. + if payload.TriggerMode == "micro" && !payload.Applied { + return transcriptDirty + } + a.state.ExecutionError = "" + a.state.StatusText = fmt.Sprintf( + "Compact(%s) saved %.1f%% context", + payload.TriggerMode, + payload.SavedRatio*100, + ) + a.appendInlineMessage( + roleSystem, + fmt.Sprintf( + "[System] Compact(%s) %s (before=%d, after=%d, saved=%.1f%%, transcript=%s)", + payload.TriggerMode, + map[bool]string{true: "applied", false: "checked"}[payload.Applied], + payload.BeforeChars, + payload.AfterChars, + payload.SavedRatio*100, + payload.TranscriptPath, + ), + ) + transcriptDirty = true + case agentruntime.EventCompactError: + payload, ok := event.Payload.(agentruntime.CompactErrorPayload) + if !ok { + return transcriptDirty + } + if payload.TriggerMode == "micro" { + // Micro compact failure is non-blocking for the main loop, keep it as a lightweight notice. + a.appendInlineMessage(roleEvent, fmt.Sprintf("Compact(micro) skipped: %s", payload.Message)) + transcriptDirty = true + return transcriptDirty + } + message := fmt.Sprintf("Compact(%s) failed: %s", payload.TriggerMode, payload.Message) + a.state.ExecutionError = message + a.state.StatusText = message + a.appendInlineMessage(roleError, message) + transcriptDirty = true } return transcriptDirty @@ -769,7 +836,7 @@ func (a *App) rebuildTranscript() { } func (a *App) handleImmediateSlashCommand(input string) (bool, tea.Cmd) { - command, _ := splitFirstWord(strings.ToLower(strings.TrimSpace(input))) + command, rest := splitFirstWord(strings.ToLower(strings.TrimSpace(input))) switch command { case slashCommandExit: return true, tea.Quit @@ -777,6 +844,18 @@ func (a *App) handleImmediateSlashCommand(input string) (bool, tea.Cmd) { a.startDraftSession() a.state.StatusText = "[System] Cleared current draft/history." return true, nil + case slashCommandCompact: + if strings.TrimSpace(rest) != "" { + errText := fmt.Sprintf("usage: %s", slashUsageCompact) + a.state.ExecutionError = errText + a.state.StatusText = errText + a.appendInlineMessage(roleError, errText) + a.rebuildTranscript() + return true, nil + } + a.state.StatusText = statusCompacting + a.state.ExecutionError = "" + return true, runCompact(a.runtime, a.state.ActiveSessionID) default: return false, nil } @@ -848,3 +927,10 @@ func runAgent(runtime agentruntime.Runtime, sessionID string, content string) te return runFinishedMsg{err: err} } } + +func runCompact(runtime agentruntime.Runtime, sessionID string) tea.Cmd { + return func() tea.Msg { + _, err := runtime.Compact(context.Background(), agentruntime.CompactInput{SessionID: sessionID}) + return compactFinishedMsg{err: err} + } +} diff --git a/internal/tui/update_test.go b/internal/tui/update_test.go index 6942e988..685490bb 100644 --- a/internal/tui/update_test.go +++ b/internal/tui/update_test.go @@ -19,15 +19,18 @@ import ( ) type stubRuntime struct { - runInputs []agentruntime.UserInput - events chan agentruntime.RuntimeEvent - sessions []agentruntime.SessionSummary - loads map[string]agentruntime.Session - runErr error - listErr error - loadErr error - cancelCalls int - cancelResult bool + runInputs []agentruntime.UserInput + compactInputs []agentruntime.CompactInput + events chan agentruntime.RuntimeEvent + sessions []agentruntime.SessionSummary + loads map[string]agentruntime.Session + runErr error + compactErr error + compactResult agentruntime.CompactResult + listErr error + loadErr error + cancelCalls int + cancelResult bool } func newStubRuntime() *stubRuntime { @@ -42,6 +45,11 @@ func (r *stubRuntime) Run(ctx context.Context, input agentruntime.UserInput) err return r.runErr } +func (r *stubRuntime) Compact(ctx context.Context, input agentruntime.CompactInput) (agentruntime.CompactResult, error) { + r.compactInputs = append(r.compactInputs, input) + return r.compactResult, r.compactErr +} + func (r *stubRuntime) Events() <-chan agentruntime.RuntimeEvent { return r.events } @@ -1191,6 +1199,93 @@ func TestAppHandleRuntimeEventAdditionalBranches(t *testing.T) { } }, }, + { + name: "compact done event emits lightweight notice", + event: agentruntime.RuntimeEvent{ + Type: agentruntime.EventCompactDone, + SessionID: "s1", + Payload: agentruntime.CompactDonePayload{ + Applied: true, + BeforeChars: 100, + AfterChars: 60, + SavedRatio: 0.4, + TriggerMode: "manual", + TranscriptPath: "/tmp/t.jsonl", + }, + }, + assert: func(t *testing.T, app App) { + t.Helper() + if !strings.Contains(app.state.StatusText, "Compact(manual)") { + t.Fatalf("expected compact status text, got %q", app.state.StatusText) + } + if len(app.activeMessages) == 0 || !strings.Contains(app.activeMessages[len(app.activeMessages)-1].Content, "Compact(manual)") { + t.Fatalf("expected compact inline notice, got %+v", app.activeMessages) + } + }, + }, + { + name: "compact done ignores passive micro checks", + setup: func(app *App) { + app.state.StatusText = "unchanged" + }, + event: agentruntime.RuntimeEvent{ + Type: agentruntime.EventCompactDone, + SessionID: "s1", + Payload: agentruntime.CompactDonePayload{ + Applied: false, + TriggerMode: "micro", + }, + }, + assert: func(t *testing.T, app App) { + t.Helper() + if app.state.StatusText != "unchanged" { + t.Fatalf("expected status unchanged for passive micro compact, got %q", app.state.StatusText) + } + }, + }, + { + name: "compact error event is surfaced", + event: agentruntime.RuntimeEvent{ + Type: agentruntime.EventCompactError, + SessionID: "s1", + Payload: agentruntime.CompactErrorPayload{ + TriggerMode: "manual", + Message: "disk full", + }, + }, + assert: func(t *testing.T, app App) { + t.Helper() + if !strings.Contains(app.state.ExecutionError, "Compact(manual) failed: disk full") { + t.Fatalf("expected compact error in state, got %q", app.state.ExecutionError) + } + }, + }, + { + name: "micro compact error remains non-blocking notice", + setup: func(app *App) { + app.state.StatusText = statusThinking + }, + event: agentruntime.RuntimeEvent{ + Type: agentruntime.EventCompactError, + SessionID: "s1", + Payload: agentruntime.CompactErrorPayload{ + TriggerMode: "micro", + Message: "tmp write failed", + }, + }, + assert: func(t *testing.T, app App) { + t.Helper() + if app.state.ExecutionError != "" { + t.Fatalf("expected no execution error for micro compact failure, got %q", app.state.ExecutionError) + } + if app.state.StatusText != statusThinking { + t.Fatalf("expected status unchanged for micro compact failure, got %q", app.state.StatusText) + } + if len(app.activeMessages) == 0 || !strings.Contains(app.activeMessages[len(app.activeMessages)-1].Content, "Compact(micro) skipped") { + t.Fatalf("expected lightweight micro compact notice, got %+v", app.activeMessages) + } + }, + }, } for _, tt := range tests { @@ -1263,6 +1358,40 @@ func TestImmediateSlashCommandsAndLayoutBranches(t *testing.T) { t.Fatalf("expected /clear to reset draft state") } + runtime.compactResult = agentruntime.CompactResult{ + Applied: true, + BeforeChars: 100, + AfterChars: 40, + SavedRatio: 0.6, + TranscriptPath: "/tmp/transcript.jsonl", + } + app.state.ActiveSessionID = "session-compact" + handled, cmd = app.handleImmediateSlashCommand("/compact") + if !handled || cmd == nil { + t.Fatalf("expected /compact to trigger compact cmd") + } + if app.state.StatusText != statusCompacting { + t.Fatalf("expected compact status %q, got %q", statusCompacting, app.state.StatusText) + } + msgs := collectTeaMessages(cmd) + if len(msgs) != 1 { + t.Fatalf("expected one compact command message, got %d", len(msgs)) + } + if _, ok := msgs[0].(compactFinishedMsg); !ok { + t.Fatalf("expected compact finished msg, got %T", msgs[0]) + } + if len(runtime.compactInputs) != 1 || runtime.compactInputs[0].SessionID != "session-compact" { + t.Fatalf("expected runtime compact call with active session, got %+v", runtime.compactInputs) + } + + handled, cmd = app.handleImmediateSlashCommand("/compact now") + if !handled || cmd != nil { + t.Fatalf("expected /compact with args to be handled locally with usage error") + } + if !strings.Contains(app.state.StatusText, "usage: /compact") { + t.Fatalf("expected /compact usage hint, got %q", app.state.StatusText) + } + handled, cmd = app.handleImmediateSlashCommand("/exit") if !handled || cmd == nil { t.Fatalf("expected /exit to return a quit cmd") From d540416d88a508aed0343fe5fb50e7cfaa0898e6 Mon Sep 17 00:00:00 2001 From: Yumiue <229866007@qq.com> Date: Fri, 3 Apr 2026 08:49:30 +0800 Subject: [PATCH 03/23] test(runtime): stabilize api key env in parallel tests --- internal/runtime/runtime_test.go | 31 +++++++++++++++---------------- 1 file changed, 15 insertions(+), 16 deletions(-) diff --git a/internal/runtime/runtime_test.go b/internal/runtime/runtime_test.go index 9ceba70e..f76ea735 100644 --- a/internal/runtime/runtime_test.go +++ b/internal/runtime/runtime_test.go @@ -17,6 +17,21 @@ import ( "neo-code/internal/tools" ) +func TestMain(m *testing.M) { + const key = config.OpenAIDefaultAPIKeyEnv + previousValue, hadPrevious := os.LookupEnv(key) + _ = os.Setenv(key, "test-key") + + exitCode := m.Run() + + if hadPrevious { + _ = os.Setenv(key, previousValue) + } else { + _ = os.Unsetenv(key) + } + os.Exit(exitCode) +} + type memoryStore struct { sessions map[string]Session saves int @@ -1449,10 +1464,6 @@ func TestServiceCompactManualFailureReturnsError(t *testing.T) { func newRuntimeConfigManager(t *testing.T) *config.Manager { t.Helper() - restoreRuntimeEnv(t, config.OpenAIDefaultAPIKeyEnv) - if err := os.Setenv(config.OpenAIDefaultAPIKeyEnv, "test-key"); err != nil { - t.Fatalf("set env: %v", err) - } manager := config.NewManager(config.NewLoader(t.TempDir(), builtin.DefaultConfig())) if _, err := manager.Load(context.Background()); err != nil { t.Fatalf("load config: %v", err) @@ -1468,18 +1479,6 @@ func newRuntimeConfigManager(t *testing.T) *config.Manager { return manager } -func restoreRuntimeEnv(t *testing.T, key string) { - t.Helper() - value, ok := os.LookupEnv(key) - t.Cleanup(func() { - if !ok { - _ = os.Unsetenv(key) - return - } - _ = os.Setenv(key, value) - }) -} - func onlySession(t *testing.T, store *memoryStore) Session { t.Helper() if len(store.sessions) != 1 { From 919499592b4f72ad19739a17cac847f2b0cd2573 Mon Sep 17 00:00:00 2001 From: Yumiue <229866007@qq.com> Date: Fri, 3 Apr 2026 09:00:28 +0800 Subject: [PATCH 04/23] refactor(runtime,tui): use shared compact trigger mode constants --- internal/runtime/events.go | 7 +++++++ internal/runtime/runtime.go | 15 +++++++++++++-- internal/tui/update.go | 4 ++-- internal/tui/update_test.go | 8 ++++---- 4 files changed, 26 insertions(+), 8 deletions(-) diff --git a/internal/runtime/events.go b/internal/runtime/events.go index 27cd73b7..33ffa1bc 100644 --- a/internal/runtime/events.go +++ b/internal/runtime/events.go @@ -13,6 +13,13 @@ type RuntimeEvent struct { Payload any } +const ( + // CompactTriggerModeMicro labels compact payloads produced by the micro strategy. + CompactTriggerModeMicro = "micro" + // CompactTriggerModeManual labels compact payloads produced by manual compact requests. + CompactTriggerModeManual = "manual" +) + const ( // EventUserMessage is emitted after the user input has been accepted and saved. EventUserMessage EventType = "user_message" diff --git a/internal/runtime/runtime.go b/internal/runtime/runtime.go index 1fe8cef1..4f78e3b7 100644 --- a/internal/runtime/runtime.go +++ b/internal/runtime/runtime.go @@ -471,7 +471,7 @@ func (s *Service) runCompactForSession( session.UpdatedAt = time.Now() if err := s.sessionStore.Save(ctx, &session); err != nil { s.emit(ctx, EventCompactError, runID, session.ID, CompactErrorPayload{ - TriggerMode: string(mode), + TriggerMode: compactTriggerMode(mode), Message: err.Error(), }) session.Messages = originalMessages @@ -487,7 +487,7 @@ func (s *Service) runCompactForSession( BeforeChars: result.Metrics.BeforeChars, AfterChars: result.Metrics.AfterChars, SavedRatio: result.Metrics.SavedRatio, - TriggerMode: result.Metrics.TriggerMode, + TriggerMode: compactTriggerMode(mode), TranscriptID: result.TranscriptID, TranscriptPath: result.TranscriptPath, } @@ -585,6 +585,17 @@ func providerRetryBackoff(attempt int) time.Duration { return wait } +func compactTriggerMode(mode contextcompact.Mode) string { + switch mode { + case contextcompact.ModeMicro: + return CompactTriggerModeMicro + case contextcompact.ModeManual: + return CompactTriggerModeManual + default: + return string(mode) + } +} + func (s *Service) isRunCanceled(err error) bool { return errors.Is(err, context.Canceled) } diff --git a/internal/tui/update.go b/internal/tui/update.go index 8af49eb8..f6352183 100644 --- a/internal/tui/update.go +++ b/internal/tui/update.go @@ -554,7 +554,7 @@ func (a *App) handleRuntimeEvent(event agentruntime.RuntimeEvent) bool { return transcriptDirty } // Skip passive micro checks that did not rewrite context to avoid transcript noise. - if payload.TriggerMode == "micro" && !payload.Applied { + if payload.TriggerMode == agentruntime.CompactTriggerModeMicro && !payload.Applied { return transcriptDirty } a.state.ExecutionError = "" @@ -581,7 +581,7 @@ func (a *App) handleRuntimeEvent(event agentruntime.RuntimeEvent) bool { if !ok { return transcriptDirty } - if payload.TriggerMode == "micro" { + if payload.TriggerMode == agentruntime.CompactTriggerModeMicro { // Micro compact failure is non-blocking for the main loop, keep it as a lightweight notice. a.appendInlineMessage(roleEvent, fmt.Sprintf("Compact(micro) skipped: %s", payload.Message)) transcriptDirty = true diff --git a/internal/tui/update_test.go b/internal/tui/update_test.go index 685490bb..1b0dd5ce 100644 --- a/internal/tui/update_test.go +++ b/internal/tui/update_test.go @@ -1209,7 +1209,7 @@ func TestAppHandleRuntimeEventAdditionalBranches(t *testing.T) { BeforeChars: 100, AfterChars: 60, SavedRatio: 0.4, - TriggerMode: "manual", + TriggerMode: agentruntime.CompactTriggerModeManual, TranscriptPath: "/tmp/t.jsonl", }, }, @@ -1233,7 +1233,7 @@ func TestAppHandleRuntimeEventAdditionalBranches(t *testing.T) { SessionID: "s1", Payload: agentruntime.CompactDonePayload{ Applied: false, - TriggerMode: "micro", + TriggerMode: agentruntime.CompactTriggerModeMicro, }, }, assert: func(t *testing.T, app App) { @@ -1249,7 +1249,7 @@ func TestAppHandleRuntimeEventAdditionalBranches(t *testing.T) { Type: agentruntime.EventCompactError, SessionID: "s1", Payload: agentruntime.CompactErrorPayload{ - TriggerMode: "manual", + TriggerMode: agentruntime.CompactTriggerModeManual, Message: "disk full", }, }, @@ -1269,7 +1269,7 @@ func TestAppHandleRuntimeEventAdditionalBranches(t *testing.T) { Type: agentruntime.EventCompactError, SessionID: "s1", Payload: agentruntime.CompactErrorPayload{ - TriggerMode: "micro", + TriggerMode: agentruntime.CompactTriggerModeMicro, Message: "tmp write failed", }, }, From ab64846c3b52f0edf1f58372a6a1659bdfe127a2 Mon Sep 17 00:00:00 2001 From: Yumiue <229866007@qq.com> Date: Fri, 3 Apr 2026 09:08:31 +0800 Subject: [PATCH 05/23] test(compact): add edge-case coverage for compact and runtime mode mapping --- internal/context/compact/runner_test.go | 140 ++++++++++++++++++++++++ internal/runtime/runtime_test.go | 36 ++++++ 2 files changed, 176 insertions(+) diff --git a/internal/context/compact/runner_test.go b/internal/context/compact/runner_test.go index 8b1608f5..5f1a5b7e 100644 --- a/internal/context/compact/runner_test.go +++ b/internal/context/compact/runner_test.go @@ -335,3 +335,143 @@ func TestSaveTranscriptUsesUniqueIDWithinSameTimestamp(t *testing.T) { t.Fatalf("second transcript file missing: %v", err) } } + +func TestRunRejectsUnsupportedMode(t *testing.T) { + t.Parallel() + + runner := NewRunner() + _, err := runner.Run(context.Background(), Input{ + Mode: Mode("invalid"), + SessionID: "session-invalid-mode", + Workdir: t.TempDir(), + Messages: []provider.Message{{Role: provider.RoleUser, Content: "hello"}}, + Config: config.CompactConfig{ + ManualStrategy: config.CompactManualStrategyKeepRecent, + ManualKeepRecentSpans: 6, + ToolResultKeepRecent: 3, + ToolResultPlaceholderMinChars: 100, + MaxSummaryChars: 1200, + }, + }) + if err == nil || !strings.Contains(err.Error(), "unsupported mode") { + t.Fatalf("expected unsupported mode error, got %v", err) + } +} + +func TestRunManualRejectsUnsupportedStrategy(t *testing.T) { + t.Parallel() + + runner := NewRunner() + home := t.TempDir() + runner.userHomeDir = func() (string, error) { return home, nil } + runner.randomToken = func() (string, error) { return "token0001", nil } + + _, err := runner.Run(context.Background(), Input{ + Mode: ModeManual, + SessionID: "session-invalid-strategy", + Workdir: t.TempDir(), + Messages: []provider.Message{{Role: provider.RoleUser, Content: "hello"}}, + Config: config.CompactConfig{ + ManualStrategy: "unknown_strategy", + ManualKeepRecentSpans: 6, + ToolResultKeepRecent: 3, + ToolResultPlaceholderMinChars: 100, + MaxSummaryChars: 1200, + }, + }) + if err == nil || !strings.Contains(err.Error(), "not supported") { + t.Fatalf("expected unsupported strategy error, got %v", err) + } +} + +func TestRunManualFullReplaceNoMessagesIsNoop(t *testing.T) { + t.Parallel() + + runner := NewRunner() + home := t.TempDir() + runner.userHomeDir = func() (string, error) { return home, nil } + runner.randomToken = func() (string, error) { return "token0002", nil } + + result, err := runner.Run(context.Background(), Input{ + Mode: ModeManual, + SessionID: "session-empty-full-replace", + Workdir: t.TempDir(), + Messages: nil, + Config: config.CompactConfig{ + ManualStrategy: config.CompactManualStrategyFullReplace, + ManualKeepRecentSpans: 6, + ToolResultKeepRecent: 3, + ToolResultPlaceholderMinChars: 100, + MaxSummaryChars: 1200, + }, + }) + if err != nil { + t.Fatalf("Run() error = %v", err) + } + if result.Applied { + t.Fatalf("expected no-op when full_replace runs on empty messages") + } + if len(result.Messages) != 0 { + t.Fatalf("expected no messages, got %+v", result.Messages) + } +} + +func TestSaveTranscriptHandlesHomeDirAndRenameFailures(t *testing.T) { + t.Parallel() + + t.Run("user home lookup failure", func(t *testing.T) { + t.Parallel() + + runner := NewRunner() + runner.userHomeDir = func() (string, error) { return "", errors.New("no home") } + _, _, err := runner.saveTranscript([]provider.Message{{Role: provider.RoleUser, Content: "hello"}}, "s1", t.TempDir()) + if err == nil || !strings.Contains(err.Error(), "resolve user home") { + t.Fatalf("expected home dir failure, got %v", err) + } + }) + + t.Run("rename failure removes temp file", func(t *testing.T) { + t.Parallel() + + runner := NewRunner() + home := t.TempDir() + runner.userHomeDir = func() (string, error) { return home, nil } + runner.randomToken = func() (string, error) { return "token0003", nil } + + removedPath := "" + runner.rename = func(oldPath, newPath string) error { + return errors.New("rename failed") + } + runner.remove = func(path string) error { + removedPath = path + return nil + } + + _, _, err := runner.saveTranscript( + []provider.Message{{Role: provider.RoleUser, Content: "hello"}}, + "s2", + filepath.Join(home, "workspace"), + ) + if err == nil || !strings.Contains(err.Error(), "commit transcript") { + t.Fatalf("expected rename commit failure, got %v", err) + } + if strings.TrimSpace(removedPath) == "" || !strings.HasSuffix(removedPath, ".tmp") { + t.Fatalf("expected temp file cleanup, got %q", removedPath) + } + }) +} + +func TestValidateSummaryRequiresDoneOrInProgress(t *testing.T) { + t.Parallel() + + _, err := validateSummary(strings.Join([]string{ + "[compact_summary]", + "decisions:", + "- none", + "constraints:", + "- keep safety checks", + }, "\n"), 0) + if err == nil || !strings.Contains(err.Error(), "requires done or in_progress content") { + t.Fatalf("expected summary section validation failure, got %v", err) + } +} diff --git a/internal/runtime/runtime_test.go b/internal/runtime/runtime_test.go index f76ea735..2699740e 100644 --- a/internal/runtime/runtime_test.go +++ b/internal/runtime/runtime_test.go @@ -1641,3 +1641,39 @@ func TestProviderRetryBackoff(t *testing.T) { }) } } + +func TestCompactTriggerMode(t *testing.T) { + t.Parallel() + + tests := []struct { + name string + mode contextcompact.Mode + want string + }{ + { + name: "micro mode", + mode: contextcompact.ModeMicro, + want: CompactTriggerModeMicro, + }, + { + name: "manual mode", + mode: contextcompact.ModeManual, + want: CompactTriggerModeManual, + }, + { + name: "unknown mode fallback", + mode: contextcompact.Mode("custom"), + want: "custom", + }, + } + + for _, tt := range tests { + tt := tt + t.Run(tt.name, func(t *testing.T) { + t.Parallel() + if got := compactTriggerMode(tt.mode); got != tt.want { + t.Fatalf("compactTriggerMode(%q) = %q, want %q", tt.mode, got, tt.want) + } + }) + } +} From 02f1c7b15fa5bba42ed634b30e92a38b336e0375 Mon Sep 17 00:00:00 2001 From: phantom5099 <1011668688@qq.com> Date: Thu, 2 Apr 2026 14:04:55 +0800 Subject: [PATCH 06/23] =?UTF-8?q?feat=EF=BC=88provider=EF=BC=89=EF=BC=9A?= =?UTF-8?q?=E5=A2=9E=E5=8A=A0=E4=B8=BB=E5=8A=A8=E8=8E=B7=E5=8F=96=E6=A8=A1?= =?UTF-8?q?=E5=9E=8B=E5=92=8C=E9=A2=84=E7=95=99=E6=A8=A1=E5=9E=8B=E5=85=83?= =?UTF-8?q?=E6=95=B0=E6=8D=AE?= MIME-Version: 1.0 Content-Type: text/plain; charset=UTF-8 Content-Transfer-Encoding: 8bit --- internal/app/bootstrap.go | 2 +- internal/config/builtin_providers.go | 37 +- internal/config/builtin_providers_test.go | 37 +- internal/config/config_test.go | 25 +- internal/config/model.go | 46 +- internal/config/provider_identity.go | 76 ++ internal/provider/discovery/discovery.go | 50 ++ internal/provider/discovery/discovery_test.go | 51 ++ internal/provider/errors.go | 11 + internal/provider/metadata.go | 265 +++++++ internal/provider/metadata_test.go | 73 ++ internal/provider/model_catalog_store.go | 130 ++++ internal/provider/model_catalog_store_test.go | 68 ++ internal/provider/openai/openai.go | 34 +- internal/provider/registry.go | 29 +- internal/provider/registry_test.go | 56 +- internal/provider/service.go | 273 ++++++-- internal/provider/service_test.go | 662 +++++++++--------- internal/provider/types.go | 20 +- internal/tui/update_test.go | 17 +- 20 files changed, 1457 insertions(+), 505 deletions(-) create mode 100644 internal/config/provider_identity.go create mode 100644 internal/provider/discovery/discovery.go create mode 100644 internal/provider/discovery/discovery_test.go create mode 100644 internal/provider/metadata.go create mode 100644 internal/provider/metadata_test.go create mode 100644 internal/provider/model_catalog_store.go create mode 100644 internal/provider/model_catalog_store_test.go diff --git a/internal/app/bootstrap.go b/internal/app/bootstrap.go index 3578c1da..b6e324f7 100644 --- a/internal/app/bootstrap.go +++ b/internal/app/bootstrap.go @@ -37,7 +37,7 @@ func NewProgram(ctx context.Context) (*tea.Program, error) { if err != nil { return nil, err } - providerService := provider.NewService(manager, providerRegistry) + providerService := provider.NewService(manager, providerRegistry, nil) sessionStore := agentruntime.NewSessionStore(loader.BaseDir()) runtimeSvc := agentruntime.NewWithFactory( diff --git a/internal/config/builtin_providers.go b/internal/config/builtin_providers.go index 90916086..57d1a2eb 100644 --- a/internal/config/builtin_providers.go +++ b/internal/config/builtin_providers.go @@ -22,78 +22,51 @@ const ( QiniuDefaultAPIKeyEnv = "QINIU_API_KEY" ) -var openAIModels = []string{ - OpenAIDefaultModel, - "gpt-4o", - "gpt-5.4", - "gpt-5.3-codex", -} - -var geminiModels = []string{ - GeminiDefaultModel, - "gemini-2.5-pro", - "gemini-2.0-flash", -} - -var openLLModels = []string{ - OpenLLDefaultModel, - "gpt-5.3-codex", - "gpt-5.3-turbo", -} - -var qiniuModels = []string{ - QiniuDefaultModel, - "openai/gpt-5.2-codex", -} - -// OpenAIProvider 返回 OpenAI provider 的默认配置。 +// OpenAIProvider returns the builtin OpenAI provider definition. func OpenAIProvider() ProviderConfig { return ProviderConfig{ Name: OpenAIName, Driver: "openai", BaseURL: OpenAIDefaultBaseURL, Model: OpenAIDefaultModel, - Models: append([]string(nil), openAIModels...), APIKeyEnv: OpenAIDefaultAPIKeyEnv, } } -// GeminiProvider 返回 Gemini provider 的默认配置。 +// GeminiProvider returns the builtin Gemini provider definition. func GeminiProvider() ProviderConfig { return ProviderConfig{ Name: GeminiName, Driver: "openai", BaseURL: GeminiDefaultBaseURL, Model: GeminiDefaultModel, - Models: append([]string(nil), geminiModels...), APIKeyEnv: GeminiDefaultAPIKeyEnv, } } -// OpenLLProvider 返回 OpenLL provider 的默认配置。 +// OpenLLProvider returns the builtin OpenLL provider definition. func OpenLLProvider() ProviderConfig { return ProviderConfig{ Name: OpenLLName, Driver: "openai", BaseURL: OpenLLDefaultBaseURL, Model: OpenLLDefaultModel, - Models: append([]string(nil), openLLModels...), APIKeyEnv: OpenLLDefaultAPIKeyEnv, } } +// QiniuProvider returns the builtin Qiniu provider definition. func QiniuProvider() ProviderConfig { return ProviderConfig{ Name: QiniuName, Driver: "openai", BaseURL: QiniuDefaultBaseURL, Model: QiniuDefaultModel, - Models: append([]string(nil), qiniuModels...), APIKeyEnv: QiniuDefaultAPIKeyEnv, } } -// DefaultProviders 返回所有内置 provider 配置列表。 +// DefaultProviders returns all builtin provider definitions. func DefaultProviders() []ProviderConfig { return []ProviderConfig{ OpenAIProvider(), diff --git a/internal/config/builtin_providers_test.go b/internal/config/builtin_providers_test.go index 5234e4f7..8e39f9a9 100644 --- a/internal/config/builtin_providers_test.go +++ b/internal/config/builtin_providers_test.go @@ -40,11 +40,11 @@ func TestOpenAIProviderConfig(t *testing.T) { if provider.APIKeyEnv != OpenAIDefaultAPIKeyEnv { t.Fatalf("expected API key env %q, got %q", OpenAIDefaultAPIKeyEnv, provider.APIKeyEnv) } - if len(provider.Models) == 0 { - t.Fatal("expected non-empty models list") + if len(provider.Models) != 0 { + t.Fatalf("expected builtin provider to leave manual models empty, got %+v", provider.Models) } - if !ContainsModelID(provider.Models, OpenAIDefaultModel) { - t.Fatalf("expected models to contain default model %q", OpenAIDefaultModel) + if !ContainsModelID(provider.SupportedModels(), OpenAIDefaultModel) { + t.Fatalf("expected supported models to contain default model %q", OpenAIDefaultModel) } } @@ -68,11 +68,11 @@ func TestGeminiProviderConfig(t *testing.T) { if provider.APIKeyEnv != GeminiDefaultAPIKeyEnv { t.Fatalf("expected API key env %q, got %q", GeminiDefaultAPIKeyEnv, provider.APIKeyEnv) } - if len(provider.Models) == 0 { - t.Fatal("expected non-empty models list") + if len(provider.Models) != 0 { + t.Fatalf("expected builtin provider to leave manual models empty, got %+v", provider.Models) } - if !ContainsModelID(provider.Models, GeminiDefaultModel) { - t.Fatalf("expected models to contain default model %q", GeminiDefaultModel) + if !ContainsModelID(provider.SupportedModels(), GeminiDefaultModel) { + t.Fatalf("expected supported models to contain default model %q", GeminiDefaultModel) } } @@ -96,27 +96,28 @@ func TestOpenLLProviderConfig(t *testing.T) { if provider.APIKeyEnv != OpenLLDefaultAPIKeyEnv { t.Fatalf("expected API key env %q, got %q", OpenLLDefaultAPIKeyEnv, provider.APIKeyEnv) } - if len(provider.Models) == 0 { - t.Fatal("expected non-empty models list") + if len(provider.Models) != 0 { + t.Fatalf("expected builtin provider to leave manual models empty, got %+v", provider.Models) } - if !ContainsModelID(provider.Models, OpenLLDefaultModel) { - t.Fatalf("expected models to contain default model %q", OpenLLDefaultModel) + if !ContainsModelID(provider.SupportedModels(), OpenLLDefaultModel) { + t.Fatalf("expected supported models to contain default model %q", OpenLLDefaultModel) } } func TestProviderModelsAreImmutable(t *testing.T) { t.Parallel() - // Verify that modifying returned models slice doesn't affect future calls provider1 := OpenAIProvider() - provider1.Models[0] = "modified-model" + models1 := provider1.SupportedModels() + models1[0] = "modified-model" provider2 := OpenAIProvider() - if provider2.Models[0] == "modified-model" { - t.Fatal("expected models slice to be independent between calls") + models2 := provider2.SupportedModels() + if models2[0] == "modified-model" { + t.Fatal("expected supported models slice to be independent between calls") } - if provider2.Models[0] != OpenAIDefaultModel { - t.Fatalf("expected first model %q, got %q", OpenAIDefaultModel, provider2.Models[0]) + if models2[0] != OpenAIDefaultModel { + t.Fatalf("expected first model %q, got %q", OpenAIDefaultModel, models2[0]) } } diff --git a/internal/config/config_test.go b/internal/config/config_test.go index 70fa4761..1d6b2937 100644 --- a/internal/config/config_test.go +++ b/internal/config/config_test.go @@ -434,13 +434,19 @@ func TestConfigValidateFailures(t *testing.T) { expectErr: "supported_content_types[0] is empty", }, { - name: "current model unsupported by selected provider", + name: "duplicate provider endpoints after normalization", config: func() *Config { cfg := validConfig.Clone() - cfg.CurrentModel = "unsupported-model" + cfg.Providers = append(cfg.Providers, ProviderConfig{ + Name: "openai-shadow", + Driver: "OPENAI", + BaseURL: "https://API.OPENAI.COM/v1/", + Model: "shadow-model", + APIKeyEnv: "OPENAI_SHADOW_KEY", + }) return &cfg }(), - expectErr: "current_model \"unsupported-model\" is not supported", + expectErr: "duplicate provider endpoint", }, } @@ -672,6 +678,19 @@ func TestApplyProviderDefaultsMatchesByNameOnly(t *testing.T) { } } +func TestApplyDefaultsPreservesDynamicCurrentModel(t *testing.T) { + t.Parallel() + + cfg := testDefaultConfig() + cfg.CurrentModel = "server-discovered-model" + + cfg.ApplyDefaultsFrom(*testDefaultConfig()) + + if cfg.CurrentModel != "server-discovered-model" { + t.Fatalf("expected dynamic current model to be preserved, got %q", cfg.CurrentModel) + } +} + func TestNormalizeWorkdirAndClone(t *testing.T) { t.Parallel() diff --git a/internal/config/model.go b/internal/config/model.go index 274ba096..3bf19fc9 100644 --- a/internal/config/model.go +++ b/internal/config/model.go @@ -133,9 +133,6 @@ func (c *Config) ApplyDefaultsFrom(defaults Config) { c.CurrentModel = defaults.CurrentModel } } - if selected, err := c.SelectedProviderConfig(); err == nil && !ContainsModelID(selected.SupportedModels(), c.CurrentModel) { - c.CurrentModel = selected.Model - } if strings.TrimSpace(c.Workdir) == "" { c.Workdir = defaults.Workdir } @@ -163,16 +160,31 @@ func (c *Config) Validate() error { } seen := make(map[string]struct{}, len(c.Providers)) + seenEndpoints := make(map[string]string, len(c.Providers)) for i, provider := range c.Providers { if err := provider.Validate(); err != nil { return fmt.Errorf("config: provider[%d]: %w", i, err) } - key := strings.ToLower(strings.TrimSpace(provider.Name)) + key := NormalizeProviderName(provider.Name) if _, exists := seen[key]; exists { return fmt.Errorf("config: duplicate provider name %q", provider.Name) } seen[key] = struct{}{} + + identity, err := provider.Identity() + if err != nil { + return fmt.Errorf("config: provider[%d]: %w", i, err) + } + if existingName, exists := seenEndpoints[identity.Key()]; exists { + return fmt.Errorf( + "config: duplicate provider endpoint %q for providers %q and %q", + identity.BaseURL, + existingName, + provider.Name, + ) + } + seenEndpoints[identity.Key()] = provider.Name } if strings.TrimSpace(c.SelectedProvider) == "" { @@ -194,9 +206,6 @@ func (c *Config) Validate() error { if strings.TrimSpace(selected.Model) == "" { return fmt.Errorf("config: selected provider %q has empty model", selected.Name) } - if !ContainsModelID(selected.SupportedModels(), c.CurrentModel) { - return fmt.Errorf("config: current_model %q is not supported by provider %q", c.CurrentModel, selected.Name) - } if err := c.Tools.Validate(); err != nil { return fmt.Errorf("config: tools: %w", err) } @@ -219,9 +228,9 @@ func (c *Config) ProviderByName(name string) (ProviderConfig, error) { return ProviderConfig{}, errors.New("config: config is nil") } - target := strings.ToLower(strings.TrimSpace(name)) + target := NormalizeProviderName(name) for _, provider := range c.Providers { - if strings.ToLower(strings.TrimSpace(provider.Name)) == target { + if NormalizeProviderName(provider.Name) == target { return provider, nil } } @@ -245,6 +254,9 @@ func (p ProviderConfig) Validate() error { if strings.TrimSpace(p.APIKeyEnv) == "" { return fmt.Errorf("provider %q api_key_env is empty", p.Name) } + if _, err := p.Identity(); err != nil { + return fmt.Errorf("provider %q: %w", p.Name, err) + } if models := normalizeModelIDs(p.Models); len(p.Models) > 0 { if len(models) == 0 { return fmt.Errorf("provider %q models is empty", p.Name) @@ -256,8 +268,16 @@ func (p ProviderConfig) Validate() error { return nil } +func (p ProviderConfig) Identity() (ProviderIdentity, error) { + return NewProviderIdentity(p.Driver, p.BaseURL) +} + +func (p ProviderConfig) ConfiguredModels() []string { + return cloneModelIDs(p.Models) +} + func (p ProviderConfig) SupportedModels() []string { - models := normalizeModelIDs(p.Models) + models := p.ConfiguredModels() if len(models) > 0 { return models } @@ -332,13 +352,13 @@ func mergeProviderDefaults(provider ProviderConfig, defaults []ProviderConfig) P } func matchDefaultProvider(provider ProviderConfig, defaults []ProviderConfig) (ProviderConfig, bool) { - name := strings.ToLower(strings.TrimSpace(provider.Name)) + name := NormalizeProviderName(provider.Name) if name == "" { return ProviderConfig{}, false } for _, candidate := range defaults { - if strings.ToLower(candidate.Name) == name { + if NormalizeProviderName(candidate.Name) == name { return candidate, true } } @@ -597,7 +617,7 @@ func normalizeModelIDs(models []string) []string { if id == "" { continue } - key := strings.ToLower(id) + key := NormalizeKey(id) if _, exists := seen[key]; exists { continue } diff --git a/internal/config/provider_identity.go b/internal/config/provider_identity.go new file mode 100644 index 00000000..93f01207 --- /dev/null +++ b/internal/config/provider_identity.go @@ -0,0 +1,76 @@ +package config + +import ( + "fmt" + "net/url" + "path" + "strings" +) + +// ProviderIdentity identifies a concrete provider endpoint for discovery and caching. +type ProviderIdentity struct { + Driver string `json:"driver"` + BaseURL string `json:"base_url"` +} + +func (i ProviderIdentity) Key() string { + return i.Driver + "|" + i.BaseURL +} + +func NormalizeKey(s string) string { + return strings.ToLower(strings.TrimSpace(s)) +} + +func NormalizeProviderName(name string) string { + return NormalizeKey(name) +} + +func NormalizeProviderDriver(driver string) string { + return NormalizeKey(driver) +} + +func NormalizeProviderBaseURL(raw string) (string, error) { + value := strings.TrimSpace(raw) + if value == "" { + return "", fmt.Errorf("provider base_url is empty") + } + + parsed, err := url.Parse(value) + if err != nil { + return "", fmt.Errorf("provider base_url %q is invalid: %w", raw, err) + } + if parsed.Scheme == "" || parsed.Host == "" { + return "", fmt.Errorf("provider base_url %q must include scheme and host", raw) + } + + parsed.Scheme = strings.ToLower(parsed.Scheme) + parsed.Host = strings.ToLower(parsed.Host) + parsed.RawQuery = "" + parsed.Fragment = "" + parsed.User = nil + + if cleaned := path.Clean(strings.TrimSpace(parsed.Path)); cleaned == "." { + parsed.Path = "" + } else { + parsed.Path = strings.TrimRight(cleaned, "/") + } + + return parsed.String(), nil +} + +func NewProviderIdentity(driver string, baseURL string) (ProviderIdentity, error) { + normalizedDriver := NormalizeProviderDriver(driver) + if normalizedDriver == "" { + return ProviderIdentity{}, fmt.Errorf("provider driver is empty") + } + + normalizedBaseURL, err := NormalizeProviderBaseURL(baseURL) + if err != nil { + return ProviderIdentity{}, err + } + + return ProviderIdentity{ + Driver: normalizedDriver, + BaseURL: normalizedBaseURL, + }, nil +} diff --git a/internal/provider/discovery/discovery.go b/internal/provider/discovery/discovery.go new file mode 100644 index 00000000..4311117c --- /dev/null +++ b/internal/provider/discovery/discovery.go @@ -0,0 +1,50 @@ +package discovery + +import ( + "context" + "encoding/json" + "fmt" + "io" + "net/http" + "strings" +) + +type openAIModelsResponse struct { + Data []map[string]any `json:"data"` +} + +// FetchOpenAICompatibleModels fetches raw model objects from an OpenAI-compatible /models endpoint. +func FetchOpenAICompatibleModels(ctx context.Context, client *http.Client, baseURL string, apiKey string) ([]map[string]any, error) { + endpoint := strings.TrimRight(strings.TrimSpace(baseURL), "/") + "/models" + req, err := http.NewRequestWithContext(ctx, http.MethodGet, endpoint, nil) + if err != nil { + return nil, fmt.Errorf("provider discovery: build request: %w", err) + } + req.Header.Set("Accept", "application/json") + if strings.TrimSpace(apiKey) != "" { + req.Header.Set("Authorization", "Bearer "+strings.TrimSpace(apiKey)) + } + + resp, err := client.Do(req) + if err != nil { + return nil, fmt.Errorf("provider discovery: send request: %w", err) + } + defer func(body io.ReadCloser) { + _ = body.Close() + }(resp.Body) + + if resp.StatusCode >= http.StatusBadRequest { + data, _ := io.ReadAll(resp.Body) + body := strings.TrimSpace(string(data)) + if body == "" { + body = resp.Status + } + return nil, fmt.Errorf("provider discovery: %s", body) + } + + var payload openAIModelsResponse + if err := json.NewDecoder(resp.Body).Decode(&payload); err != nil { + return nil, fmt.Errorf("provider discovery: decode response: %w", err) + } + return payload.Data, nil +} diff --git a/internal/provider/discovery/discovery_test.go b/internal/provider/discovery/discovery_test.go new file mode 100644 index 00000000..50380a4c --- /dev/null +++ b/internal/provider/discovery/discovery_test.go @@ -0,0 +1,51 @@ +package discovery + +import ( + "context" + "net/http" + "net/http/httptest" + "strings" + "testing" +) + +func TestFetchOpenAICompatibleModels(t *testing.T) { + t.Parallel() + + server := httptest.NewServer(http.HandlerFunc(func(w http.ResponseWriter, r *http.Request) { + if r.URL.Path != "/models" { + t.Fatalf("unexpected path %q", r.URL.Path) + } + if got := r.Header.Get("Authorization"); got != "Bearer test-key" { + t.Fatalf("unexpected auth header %q", got) + } + w.Header().Set("Content-Type", "application/json") + _, _ = w.Write([]byte(`{"data":[{"id":"gpt-test","context_window":128000,"extra":"kept"}]}`)) + })) + defer server.Close() + + models, err := FetchOpenAICompatibleModels(context.Background(), server.Client(), server.URL, "test-key") + if err != nil { + t.Fatalf("FetchOpenAICompatibleModels() error = %v", err) + } + if len(models) != 1 || models[0]["id"] != "gpt-test" { + t.Fatalf("unexpected models payload: %+v", models) + } + if models[0]["extra"] != "kept" { + t.Fatalf("expected unknown fields to remain in raw payload, got %+v", models[0]) + } +} + +func TestFetchOpenAICompatibleModelsHTTPError(t *testing.T) { + t.Parallel() + + server := httptest.NewServer(http.HandlerFunc(func(w http.ResponseWriter, r *http.Request) { + w.WriteHeader(http.StatusBadGateway) + _, _ = w.Write([]byte("gateway failed")) + })) + defer server.Close() + + _, err := FetchOpenAICompatibleModels(context.Background(), server.Client(), server.URL, "test-key") + if err == nil || !strings.Contains(err.Error(), "gateway failed") { + t.Fatalf("expected gateway failure error, got %v", err) + } +} diff --git a/internal/provider/errors.go b/internal/provider/errors.go index 083fb39f..2280487a 100644 --- a/internal/provider/errors.go +++ b/internal/provider/errors.go @@ -1,10 +1,21 @@ package provider import ( + "errors" "fmt" "net/http" ) +// 通用领域错误。 +var ( + ErrProviderNotFound = errors.New("provider not found") + ErrModelNotFound = errors.New("model not found") + ErrDriverNotFound = errors.New("provider driver not found") + ErrDriverAlreadyRegistered = errors.New("provider: driver already registered") + ErrServiceManagerNil = errors.New("provider: config manager is nil") + ErrServiceRegistryNil = errors.New("provider: registry is nil") +) + type ProviderErrorCode string const ( diff --git a/internal/provider/metadata.go b/internal/provider/metadata.go new file mode 100644 index 00000000..cf04de01 --- /dev/null +++ b/internal/provider/metadata.go @@ -0,0 +1,265 @@ +package provider + +import ( + "strings" + + "neo-code/internal/config" +) + +// DescriptorFromMetadata normalizes a raw provider model object into a ModelDescriptor. +func DescriptorFromMetadata(raw map[string]any) (ModelDescriptor, bool) { + id := firstNonEmptyString( + stringValue(raw["id"]), + stringValue(raw["model"]), + stringValue(raw["name"]), + ) + if id == "" { + return ModelDescriptor{}, false + } + + descriptor := ModelDescriptor{ + ID: id, + Name: firstNonEmptyString(stringValue(raw["name"]), stringValue(raw["display_name"]), id), + Description: stringValue(raw["description"]), + ContextWindow: firstPositiveInt(raw["context_window"], raw["contextLength"], raw["input_token_limit"], raw["max_context_tokens"]), + MaxOutputTokens: firstPositiveInt(raw["max_output_tokens"], raw["output_token_limit"], raw["max_tokens"]), + Capabilities: boolMapValue(raw["capabilities"]), + Metadata: cloneStringAnyMap(raw), + } + return normalizeModelDescriptor(descriptor), true +} + +func modelDescriptorsFromIDs(modelIDs []string) []ModelDescriptor { + if len(modelIDs) == 0 { + return nil + } + + descriptors := make([]ModelDescriptor, 0, len(modelIDs)) + for _, modelID := range modelIDs { + id := strings.TrimSpace(modelID) + if id == "" { + continue + } + descriptors = append(descriptors, ModelDescriptor{ + ID: id, + Name: id, + }) + } + if len(descriptors) == 0 { + return nil + } + return descriptors +} + +func MergeModelDescriptors(sources ...[]ModelDescriptor) []ModelDescriptor { + if len(sources) == 0 { + return nil + } + + merged := make([]ModelDescriptor, 0) + indexByID := make(map[string]int) + + for _, source := range sources { + for _, candidate := range source { + normalized := normalizeModelDescriptor(candidate) + key := normalizeModelID(normalized.ID) + if key == "" { + continue + } + + if index, exists := indexByID[key]; exists { + merged[index] = mergeModelDescriptor(merged[index], normalized) + continue + } + + indexByID[key] = len(merged) + merged = append(merged, normalized) + } + } + + if len(merged) == 0 { + return nil + } + return merged +} + +func normalizeModelDescriptor(descriptor ModelDescriptor) ModelDescriptor { + descriptor.ID = strings.TrimSpace(descriptor.ID) + descriptor.Name = strings.TrimSpace(descriptor.Name) + descriptor.Description = strings.TrimSpace(descriptor.Description) + if descriptor.Name == "" { + descriptor.Name = descriptor.ID + } + descriptor.Capabilities = cloneStringBoolMap(descriptor.Capabilities) + descriptor.Metadata = cloneStringAnyMap(descriptor.Metadata) + return descriptor +} + +func mergeModelDescriptor(primary ModelDescriptor, secondary ModelDescriptor) ModelDescriptor { + if strings.TrimSpace(primary.Name) == "" { + primary.Name = secondary.Name + } + if strings.TrimSpace(primary.Description) == "" { + primary.Description = secondary.Description + } + if primary.ContextWindow <= 0 { + primary.ContextWindow = secondary.ContextWindow + } + if primary.MaxOutputTokens <= 0 { + primary.MaxOutputTokens = secondary.MaxOutputTokens + } + primary.Capabilities = mergeStringBoolMaps(primary.Capabilities, secondary.Capabilities) + primary.Metadata = mergeStringAnyMaps(primary.Metadata, secondary.Metadata) + return normalizeModelDescriptor(primary) +} + +func mergeStringBoolMaps(primary map[string]bool, secondary map[string]bool) map[string]bool { + if len(primary) == 0 && len(secondary) == 0 { + return nil + } + + merged := cloneStringBoolMap(primary) + if merged == nil { + merged = make(map[string]bool, len(secondary)) + } + for key, value := range secondary { + if _, exists := merged[key]; exists { + continue + } + merged[key] = value + } + return merged +} + +func mergeStringAnyMaps(primary map[string]any, secondary map[string]any) map[string]any { + if len(primary) == 0 && len(secondary) == 0 { + return nil + } + + merged := cloneStringAnyMap(primary) + if merged == nil { + merged = make(map[string]any, len(secondary)) + } + + for key, value := range secondary { + current, exists := merged[key] + if !exists { + merged[key] = cloneJSONValue(value) + continue + } + + currentMap, currentOK := current.(map[string]any) + valueMap, valueOK := value.(map[string]any) + if currentOK && valueOK { + merged[key] = mergeStringAnyMaps(currentMap, valueMap) + } + } + return merged +} + +func boolMapValue(value any) map[string]bool { + raw, ok := value.(map[string]any) + if !ok || len(raw) == 0 { + return nil + } + + normalized := make(map[string]bool, len(raw)) + for key, value := range raw { + boolValue, ok := value.(bool) + if !ok { + continue + } + normalized[strings.TrimSpace(key)] = boolValue + } + if len(normalized) == 0 { + return nil + } + return normalized +} + +func stringValue(value any) string { + switch typed := value.(type) { + case string: + return strings.TrimSpace(typed) + default: + return "" + } +} + +func firstNonEmptyString(values ...string) string { + for _, value := range values { + if strings.TrimSpace(value) != "" { + return strings.TrimSpace(value) + } + } + return "" +} + +func firstPositiveInt(values ...any) int { + for _, value := range values { + switch typed := value.(type) { + case int: + if typed > 0 { + return typed + } + case int32: + if typed > 0 { + return int(typed) + } + case int64: + if typed > 0 { + return int(typed) + } + case float32: + if typed > 0 { + return int(typed) + } + case float64: + if typed > 0 { + return int(typed) + } + } + } + return 0 +} + +func cloneStringBoolMap(source map[string]bool) map[string]bool { + if len(source) == 0 { + return nil + } + cloned := make(map[string]bool, len(source)) + for key, value := range source { + cloned[key] = value + } + return cloned +} + +func cloneStringAnyMap(source map[string]any) map[string]any { + if len(source) == 0 { + return nil + } + cloned := make(map[string]any, len(source)) + for key, value := range source { + cloned[key] = cloneJSONValue(value) + } + return cloned +} + +func cloneJSONValue(value any) any { + switch typed := value.(type) { + case map[string]any: + return cloneStringAnyMap(typed) + case []any: + cloned := make([]any, 0, len(typed)) + for _, item := range typed { + cloned = append(cloned, cloneJSONValue(item)) + } + return cloned + default: + return typed + } +} + +func normalizeModelID(id string) string { + return config.NormalizeKey(id) +} diff --git a/internal/provider/metadata_test.go b/internal/provider/metadata_test.go new file mode 100644 index 00000000..d1033fc1 --- /dev/null +++ b/internal/provider/metadata_test.go @@ -0,0 +1,73 @@ +package provider + +import "testing" + +func TestDescriptorFromMetadataPreservesRawFields(t *testing.T) { + t.Parallel() + + raw := map[string]any{ + "id": "gpt-test", + "name": "GPT Test", + "description": "custom model", + "context_window": float64(128000), + "max_output_tokens": float64(8192), + "capabilities": map[string]any{ + "tool_call": true, + "vision": false, + "notes": "ignored", + }, + "experimental": map[string]any{ + "tier": "beta", + }, + } + + descriptor, ok := DescriptorFromMetadata(raw) + if !ok { + t.Fatal("expected descriptor to be normalized") + } + if descriptor.ID != "gpt-test" || descriptor.Name != "GPT Test" { + t.Fatalf("unexpected descriptor identity: %+v", descriptor) + } + if descriptor.ContextWindow != 128000 || descriptor.MaxOutputTokens != 8192 { + t.Fatalf("expected token metadata to be normalized, got %+v", descriptor) + } + if !descriptor.Capabilities["tool_call"] { + t.Fatalf("expected tool_call capability, got %+v", descriptor.Capabilities) + } + if _, ok := descriptor.Metadata["experimental"]; !ok { + t.Fatalf("expected unknown metadata to be preserved, got %+v", descriptor.Metadata) + } +} + +func TestMergeModelDescriptorsPrefersEarlierSourceAndBackfillsMetadata(t *testing.T) { + t.Parallel() + + manual := []ModelDescriptor{{ + ID: "gpt-test", + Name: "Manual Name", + }} + discovered := []ModelDescriptor{{ + ID: "gpt-test", + Name: "Discovered Name", + ContextWindow: 64000, + MaxOutputTokens: 4096, + Metadata: map[string]any{ + "id": "gpt-test", + "context_window": float64(64000), + }, + }} + + merged := MergeModelDescriptors(manual, discovered) + if len(merged) != 1 { + t.Fatalf("expected 1 merged model, got %d", len(merged)) + } + if merged[0].Name != "Manual Name" { + t.Fatalf("expected earlier source to win for name, got %+v", merged[0]) + } + if merged[0].ContextWindow != 64000 { + t.Fatalf("expected metadata to be backfilled, got %+v", merged[0]) + } + if _, ok := merged[0].Metadata["context_window"]; !ok { + t.Fatalf("expected raw metadata to survive merge, got %+v", merged[0].Metadata) + } +} diff --git a/internal/provider/model_catalog_store.go b/internal/provider/model_catalog_store.go new file mode 100644 index 00000000..bca392b9 --- /dev/null +++ b/internal/provider/model_catalog_store.go @@ -0,0 +1,130 @@ +package provider + +import ( + "context" + "crypto/sha256" + "encoding/hex" + "encoding/json" + "errors" + "fmt" + "os" + "path/filepath" + "strings" + "time" + + "neo-code/internal/config" +) + +const modelCatalogSchemaVersion = 1 + +var ErrModelCatalogNotFound = errors.New("provider: model catalog not found") + +// ModelCatalog stores discovered provider models for a concrete endpoint. +type ModelCatalog struct { + SchemaVersion int `json:"schema_version"` + Identity config.ProviderIdentity `json:"identity"` + FetchedAt time.Time `json:"fetched_at"` + ExpiresAt time.Time `json:"expires_at"` + Models []ModelDescriptor `json:"models"` +} + +func (c ModelCatalog) Expired(now time.Time) bool { + return !c.ExpiresAt.IsZero() && !now.Before(c.ExpiresAt) +} + +// ModelCatalogStore persists model catalogs keyed by normalized provider identity. +type ModelCatalogStore interface { + Load(ctx context.Context, identity config.ProviderIdentity) (ModelCatalog, error) + Save(ctx context.Context, catalog ModelCatalog) error +} + +type JSONModelCatalogStore struct { + dir string +} + +func NewJSONModelCatalogStore(baseDir string) *JSONModelCatalogStore { + return &JSONModelCatalogStore{ + dir: filepath.Join(strings.TrimSpace(baseDir), "cache", "models"), + } +} + +func (s *JSONModelCatalogStore) Load(ctx context.Context, identity config.ProviderIdentity) (ModelCatalog, error) { + if err := ctx.Err(); err != nil { + return ModelCatalog{}, err + } + + normalized, err := config.NewProviderIdentity(identity.Driver, identity.BaseURL) + if err != nil { + return ModelCatalog{}, fmt.Errorf("provider: normalize model catalog key: %w", err) + } + + path, err := s.catalogPathFromNormalized(normalized) + if err != nil { + return ModelCatalog{}, err + } + + data, err := os.ReadFile(path) + if err != nil { + if os.IsNotExist(err) { + return ModelCatalog{}, ErrModelCatalogNotFound + } + return ModelCatalog{}, fmt.Errorf("provider: read model catalog: %w", err) + } + + var catalog ModelCatalog + if err := json.Unmarshal(data, &catalog); err != nil { + return ModelCatalog{}, fmt.Errorf("provider: decode model catalog: %w", err) + } + + catalog.Identity = normalized + return normalizeModelCatalog(catalog), nil +} + +func (s *JSONModelCatalogStore) Save(ctx context.Context, catalog ModelCatalog) error { + if err := ctx.Err(); err != nil { + return err + } + + normalized := normalizeModelCatalog(catalog) + + normalizedIdentity, err := config.NewProviderIdentity(normalized.Identity.Driver, normalized.Identity.BaseURL) + if err != nil { + return fmt.Errorf("provider: normalize model catalog key: %w", err) + } + normalized.Identity = normalizedIdentity + + path, err := s.catalogPathFromNormalized(normalized.Identity) + if err != nil { + return err + } + + if err := os.MkdirAll(s.dir, 0o755); err != nil { + return fmt.Errorf("provider: create model catalog dir: %w", err) + } + + data, err := json.MarshalIndent(normalized, "", " ") + if err != nil { + return fmt.Errorf("provider: encode model catalog: %w", err) + } + if len(data) == 0 || data[len(data)-1] != '\n' { + data = append(data, '\n') + } + + if err := os.WriteFile(path, data, 0o644); err != nil { + return fmt.Errorf("provider: write model catalog: %w", err) + } + return nil +} + +func normalizeModelCatalog(catalog ModelCatalog) ModelCatalog { + if catalog.SchemaVersion == 0 { + catalog.SchemaVersion = modelCatalogSchemaVersion + } + catalog.Models = MergeModelDescriptors(catalog.Models) + return catalog +} + +func (s *JSONModelCatalogStore) catalogPathFromNormalized(identity config.ProviderIdentity) (string, error) { + sum := sha256.Sum256([]byte(identity.Key())) + return filepath.Join(s.dir, hex.EncodeToString(sum[:])+".json"), nil +} diff --git a/internal/provider/model_catalog_store_test.go b/internal/provider/model_catalog_store_test.go new file mode 100644 index 00000000..73f55396 --- /dev/null +++ b/internal/provider/model_catalog_store_test.go @@ -0,0 +1,68 @@ +package provider + +import ( + "context" + "errors" + "testing" + "time" + + "neo-code/internal/config" +) + +func TestJSONModelCatalogStoreRoundTrip(t *testing.T) { + t.Parallel() + + store := NewJSONModelCatalogStore(t.TempDir()) + identity, err := config.NewProviderIdentity("OPENAI", "https://API.OPENAI.COM/v1/") + if err != nil { + t.Fatalf("NewProviderIdentity() error = %v", err) + } + + expected := ModelCatalog{ + SchemaVersion: modelCatalogSchemaVersion, + Identity: identity, + FetchedAt: time.Date(2026, 4, 2, 10, 0, 0, 0, time.UTC), + ExpiresAt: time.Date(2026, 4, 3, 10, 0, 0, 0, time.UTC), + Models: []ModelDescriptor{{ + ID: "gpt-test", + Name: "GPT Test", + Metadata: map[string]any{ + "id": "gpt-test", + "experimental": map[string]any{"tier": "beta"}, + }, + }}, + } + + if err := store.Save(context.Background(), expected); err != nil { + t.Fatalf("Save() error = %v", err) + } + + got, err := store.Load(context.Background(), identity) + if err != nil { + t.Fatalf("Load() error = %v", err) + } + if got.Identity.BaseURL != "https://api.openai.com/v1" { + t.Fatalf("expected normalized base url, got %q", got.Identity.BaseURL) + } + if len(got.Models) != 1 { + t.Fatalf("expected 1 model, got %+v", got.Models) + } + if _, ok := got.Models[0].Metadata["experimental"]; !ok { + t.Fatalf("expected nested metadata to survive round-trip, got %+v", got.Models[0].Metadata) + } +} + +func TestJSONModelCatalogStoreMissingCatalog(t *testing.T) { + t.Parallel() + + store := NewJSONModelCatalogStore(t.TempDir()) + identity, err := config.NewProviderIdentity("openai", "https://api.openai.com/v1") + if err != nil { + t.Fatalf("NewProviderIdentity() error = %v", err) + } + + _, err = store.Load(context.Background(), identity) + if !errors.Is(err, ErrModelCatalogNotFound) { + t.Fatalf("expected ErrModelCatalogNotFound, got %v", err) + } +} diff --git a/internal/provider/openai/openai.go b/internal/provider/openai/openai.go index 700b50a3..0d33348d 100644 --- a/internal/provider/openai/openai.go +++ b/internal/provider/openai/openai.go @@ -16,6 +16,7 @@ import ( "neo-code/internal/config" domain "neo-code/internal/provider" + modeldiscovery "neo-code/internal/provider/discovery" "neo-code/internal/provider/transport" ) @@ -39,14 +40,24 @@ func WithTransport(rt http.RoundTripper) BuildOption { const DriverName = "openai" +// defaultRetryTransport 返回内置的带重试的 HTTP Transport。 +func defaultRetryTransport() http.RoundTripper { + return transport.NewRetryTransport(http.DefaultTransport, transport.DefaultRetryConfig()) +} + // Driver 返回 OpenAI 协议驱动的定义。 func Driver() domain.DriverDefinition { return domain.DriverDefinition{ Name: DriverName, Build: func(ctx context.Context, cfg config.ResolvedProviderConfig) (domain.Provider, error) { - baseTransport := http.DefaultTransport - retryTransport := transport.NewRetryTransport(baseTransport, transport.DefaultRetryConfig()) - return New(cfg, WithTransport(retryTransport)) + return New(cfg, WithTransport(defaultRetryTransport())) + }, + Discover: func(ctx context.Context, cfg config.ResolvedProviderConfig) ([]domain.ModelDescriptor, error) { + provider, err := New(cfg, WithTransport(defaultRetryTransport())) + if err != nil { + return nil, err + } + return provider.DiscoverModels(ctx) }, } } @@ -75,6 +86,23 @@ func New(cfg config.ResolvedProviderConfig, opts ...BuildOption) (*Provider, err }, nil } +func (p *Provider) DiscoverModels(ctx context.Context) ([]domain.ModelDescriptor, error) { + rawModels, err := modeldiscovery.FetchOpenAICompatibleModels(ctx, p.client, p.cfg.BaseURL, p.cfg.APIKey) + if err != nil { + return nil, err + } + + descriptors := make([]domain.ModelDescriptor, 0, len(rawModels)) + for _, raw := range rawModels { + descriptor, ok := domain.DescriptorFromMetadata(raw) + if !ok { + continue + } + descriptors = append(descriptors, descriptor) + } + return domain.MergeModelDescriptors(descriptors), nil +} + func (p *Provider) Chat(ctx context.Context, req domain.ChatRequest, events chan<- domain.StreamEvent) (domain.ChatResponse, error) { payload, err := p.buildRequest(req) if err != nil { diff --git a/internal/provider/registry.go b/internal/provider/registry.go index c017b04c..f5e17324 100644 --- a/internal/provider/registry.go +++ b/internal/provider/registry.go @@ -10,18 +10,18 @@ import ( ) type Builder func(ctx context.Context, cfg config.ResolvedProviderConfig) (Provider, error) +type DiscoveryFunc func(ctx context.Context, cfg config.ResolvedProviderConfig) ([]ModelDescriptor, error) type DriverDefinition struct { - Name string - Build Builder + Name string + Build Builder + Discover DiscoveryFunc } type Registry struct { drivers map[string]DriverDefinition } -var errDriverAlreadyRegistered = errors.New("provider: driver already registered") - func NewRegistry() *Registry { return &Registry{drivers: map[string]DriverDefinition{}} } @@ -34,7 +34,7 @@ func (r *Registry) Register(driver DriverDefinition) error { r.ensureDrivers() driver.Name = strings.TrimSpace(driver.Name) - driverType := normalizeKey(driver.Name) + driverType := config.NormalizeKey(driver.Name) if driverType == "" { return errors.New("provider: driver name is empty") } @@ -42,7 +42,7 @@ func (r *Registry) Register(driver DriverDefinition) error { return fmt.Errorf("provider: driver %q build func is nil", driver.Name) } if _, exists := r.drivers[driverType]; exists { - return fmt.Errorf("%w: %s", errDriverAlreadyRegistered, driver.Name) + return fmt.Errorf("%w: %s", ErrDriverAlreadyRegistered, driver.Name) } r.drivers[driverType] = driver return nil @@ -56,6 +56,17 @@ func (r *Registry) Build(ctx context.Context, cfg config.ResolvedProviderConfig) return driver.Build(ctx, cfg) } +func (r *Registry) DiscoverModels(ctx context.Context, cfg config.ResolvedProviderConfig) ([]ModelDescriptor, error) { + driver, err := r.driver(cfg.Driver) + if err != nil { + return nil, err + } + if driver.Discover == nil { + return nil, nil + } + return driver.Discover(ctx, cfg) +} + func (r *Registry) Supports(driverType string) bool { _, err := r.driver(driverType) return err == nil @@ -65,17 +76,13 @@ func (r *Registry) driver(driverType string) (DriverDefinition, error) { if r == nil { return DriverDefinition{}, ErrDriverNotFound } - driver, ok := r.drivers[normalizeKey(driverType)] + driver, ok := r.drivers[config.NormalizeKey(driverType)] if !ok { return DriverDefinition{}, fmt.Errorf("%w: %s", ErrDriverNotFound, strings.TrimSpace(driverType)) } return driver, nil } -func normalizeKey(value string) string { - return strings.ToLower(strings.TrimSpace(value)) -} - func (r *Registry) ensureDrivers() { if r.drivers == nil { r.drivers = map[string]DriverDefinition{} diff --git a/internal/provider/registry_test.go b/internal/provider/registry_test.go index 16c022ca..5627e9ca 100644 --- a/internal/provider/registry_test.go +++ b/internal/provider/registry_test.go @@ -3,6 +3,7 @@ package provider_test import ( "context" "errors" + "sync" "testing" "neo-code/internal/config" @@ -116,16 +117,16 @@ func TestServiceListProvidersUsesConfiguredMetadata(t *testing.T) { t.Fatalf("append provider: %v", err) } - service := provider.NewService(manager, newTestRegistry(t)) + service := provider.NewService(manager, newTestRegistry(t), newCatalogStoreStub()) items, err := service.ListProviders(context.Background()) if err != nil { t.Fatalf("ListProviders() error = %v", err) } expectedModels := map[string]int{ - config.OpenAIName: len(config.OpenAIProvider().Models), - config.GeminiName: len(config.GeminiProvider().Models), - config.OpenLLName: len(config.OpenLLProvider().Models), - config.QiniuName: len(config.QiniuProvider().Models), + config.OpenAIName: len(config.OpenAIProvider().SupportedModels()), + config.GeminiName: len(config.GeminiProvider().SupportedModels()), + config.OpenLLName: len(config.OpenLLProvider().SupportedModels()), + config.QiniuName: len(config.QiniuProvider().SupportedModels()), } if len(items) != len(expectedModels) { t.Fatalf("expected only supported providers, got %d", len(items)) @@ -172,7 +173,7 @@ func TestServiceSelectProviderAndSetCurrentModel(t *testing.T) { t.Fatalf("seed current model: %v", err) } - service := provider.NewService(manager, registry) + service := provider.NewService(manager, registry, newCatalogStoreStub()) selection, err := service.SelectProvider(context.Background(), "custom-main") if err != nil { @@ -239,7 +240,7 @@ func TestServiceModelOperationsUseProviderConfigEvenWithoutDriver(t *testing.T) t.Fatalf("append broken provider: %v", err) } - service := provider.NewService(manager, newTestRegistry(t)) + service := provider.NewService(manager, newTestRegistry(t), newCatalogStoreStub()) models, err := service.ListModels(context.Background()) if err != nil { @@ -276,7 +277,7 @@ func TestServiceSelectProviderRejectsUnsupportedDriver(t *testing.T) { t.Fatalf("append broken provider: %v", err) } - service := provider.NewService(manager, newTestRegistry(t)) + service := provider.NewService(manager, newTestRegistry(t), newCatalogStoreStub()) if _, err := service.SelectProvider(context.Background(), "broken-provider"); !errors.Is(err, provider.ErrDriverNotFound) { t.Fatalf("expected SelectProvider() to preserve ErrDriverNotFound, got %v", err) } @@ -323,10 +324,47 @@ func TestServiceBuildReturnsErrorOnNilRegistry(t *testing.T) { t.Parallel() manager := newTestManager(t) - service := provider.NewService(manager, nil) + service := provider.NewService(manager, nil, newCatalogStoreStub()) _, err := service.Build(context.Background(), config.ResolvedProviderConfig{}) if err == nil { t.Fatalf("expected error for nil registry") } } + +type catalogStoreStub struct { + mu sync.Mutex + catalogs map[string]provider.ModelCatalog +} + +func newCatalogStoreStub() *catalogStoreStub { + return &catalogStoreStub{ + catalogs: map[string]provider.ModelCatalog{}, + } +} + +func (s *catalogStoreStub) Load(ctx context.Context, identity config.ProviderIdentity) (provider.ModelCatalog, error) { + if err := ctx.Err(); err != nil { + return provider.ModelCatalog{}, err + } + + s.mu.Lock() + defer s.mu.Unlock() + + catalog, ok := s.catalogs[identity.Key()] + if !ok { + return provider.ModelCatalog{}, provider.ErrModelCatalogNotFound + } + return catalog, nil +} + +func (s *catalogStoreStub) Save(ctx context.Context, catalog provider.ModelCatalog) error { + if err := ctx.Err(); err != nil { + return err + } + + s.mu.Lock() + defer s.mu.Unlock() + s.catalogs[catalog.Identity.Key()] = catalog + return nil +} diff --git a/internal/provider/service.go b/internal/provider/service.go index acbac2db..e0e8b18e 100644 --- a/internal/provider/service.go +++ b/internal/provider/service.go @@ -2,26 +2,43 @@ package provider import ( "context" - "errors" "strings" + "sync" + "time" "neo-code/internal/config" ) +const ( + defaultModelCatalogTTL = 24 * time.Hour + defaultBackgroundRefreshTime = 30 * time.Second +) + type Service struct { - manager *config.Manager - registry *Registry + manager *config.Manager + registry *Registry + catalogs ModelCatalogStore + catalogTTL time.Duration + backgroundTimeout time.Duration + now func() time.Time + + refreshMu sync.Mutex + inFlightByID map[string]struct{} } -var ( - errServiceManagerNil = errors.New("provider: config manager is nil") - errServiceRegistryNil = errors.New("provider: registry is nil") -) +func NewService(manager *config.Manager, registry *Registry, catalogs ModelCatalogStore) *Service { + if catalogs == nil && manager != nil { + catalogs = NewJSONModelCatalogStore(manager.BaseDir()) + } -func NewService(manager *config.Manager, registry *Registry) *Service { return &Service{ - manager: manager, - registry: registry, + manager: manager, + registry: registry, + catalogs: catalogs, + catalogTTL: defaultModelCatalogTTL, + backgroundTimeout: defaultBackgroundRefreshTime, + now: time.Now, + inFlightByID: map[string]struct{}{}, } } @@ -42,7 +59,9 @@ func (s *Service) ListProviders(ctx context.Context) ([]ProviderCatalogItem, err if !s.registry.Supports(providerCfg.Driver) { continue } - items = append(items, catalogItemFromConfig(providerCfg)) + + models := s.modelsForProvider(ctx, providerCfg, modelQueryOptions{}) + items = append(items, catalogItemFromConfig(providerCfg, models)) } return items, nil @@ -53,19 +72,29 @@ func (s *Service) SelectProvider(ctx context.Context, providerName string) (Prov return ProviderSelection{}, err } - var selection ProviderSelection + cfgSnapshot := s.manager.Get() + providerCfg, err := cfgSnapshot.ProviderByName(providerName) + if err != nil { + return ProviderSelection{}, ErrProviderNotFound + } + if !s.registry.Supports(providerCfg.Driver) { + return ProviderSelection{}, ErrDriverNotFound + } + + models := s.modelsForProvider(ctx, providerCfg, modelQueryOptions{}) - err := s.manager.Update(ctx, func(cfg *config.Config) error { - providerCfg, err := cfg.ProviderByName(providerName) + var selection ProviderSelection + err = s.manager.Update(ctx, func(cfg *config.Config) error { + selected, err := cfg.ProviderByName(providerName) if err != nil { return ErrProviderNotFound } - if !s.registry.Supports(providerCfg.Driver) { + if !s.registry.Supports(selected.Driver) { return ErrDriverNotFound } - cfg.SelectedProvider = providerCfg.Name - cfg.CurrentModel = selectModel(cfg.CurrentModel, providerCfg.SupportedModels(), providerCfg.Model) + cfg.SelectedProvider = selected.Name + cfg.CurrentModel = selectModel(cfg.CurrentModel, modelDescriptorIDs(models), selected.Model) selection = selectionFromConfig(*cfg) return nil }) @@ -73,6 +102,7 @@ func (s *Service) SelectProvider(ctx context.Context, providerName string) (Prov return ProviderSelection{}, err } + s.queueRefresh(providerCfg) return selection, nil } @@ -90,7 +120,10 @@ func (s *Service) ListModels(ctx context.Context) ([]ModelDescriptor, error) { return nil, err } - return modelDescriptors(selected.SupportedModels()), nil + return s.modelsForProvider(ctx, selected, modelQueryOptions{ + allowSyncRefresh: true, + queueRefresh: true, + }), nil } func (s *Service) SetCurrentModel(ctx context.Context, modelID string) (ProviderSelection, error) { @@ -103,17 +136,25 @@ func (s *Service) SetCurrentModel(ctx context.Context, modelID string) (Provider return ProviderSelection{}, ErrModelNotFound } + cfgSnapshot := s.manager.Get() + selected, err := cfgSnapshot.SelectedProviderConfig() + if err != nil { + return ProviderSelection{}, err + } + + models := s.modelsForProvider(ctx, selected, modelQueryOptions{ + allowSyncRefresh: true, + queueRefresh: true, + }) + if !containsModelDescriptorID(models, modelID) { + return ProviderSelection{}, ErrModelNotFound + } + var selection ProviderSelection - err := s.manager.Update(ctx, func(cfg *config.Config) error { - selected, err := cfg.SelectedProviderConfig() - if err != nil { + err = s.manager.Update(ctx, func(cfg *config.Config) error { + if _, err := cfg.SelectedProviderConfig(); err != nil { return err } - - if !config.ContainsModelID(selected.SupportedModels(), modelID) { - return ErrModelNotFound - } - cfg.CurrentModel = modelID selection = selectionFromConfig(*cfg) return nil @@ -125,6 +166,147 @@ func (s *Service) SetCurrentModel(ctx context.Context, modelID string) (Provider return selection, nil } +func (s *Service) Build(ctx context.Context, cfg config.ResolvedProviderConfig) (Provider, error) { + if err := s.validate(); err != nil { + return nil, err + } + return s.registry.Build(ctx, cfg) +} + +func (s *Service) validate() error { + if s == nil || s.manager == nil { + return ErrServiceManagerNil + } + if s.registry == nil { + return ErrServiceRegistryNil + } + return nil +} + +type modelQueryOptions struct { + allowSyncRefresh bool + queueRefresh bool +} + +func (s *Service) modelsForProvider(ctx context.Context, providerCfg config.ProviderConfig, options modelQueryOptions) []ModelDescriptor { + manualModels := modelDescriptorsFromIDs(providerCfg.ConfiguredModels()) + defaultModels := modelDescriptorsFromIDs([]string{providerCfg.Model}) + + cached, cachedOK := s.loadCatalogModels(ctx, providerCfg) + if !cachedOK && options.allowSyncRefresh { + discovered, ok := s.discoverAndPersist(ctx, providerCfg) + if ok { + cached = discovered + cachedOK = true + } + } + + if cachedOK && options.queueRefresh && s.catalogExpired(ctx, providerCfg) { + s.queueRefresh(providerCfg) + } + + return MergeModelDescriptors(manualModels, cached, defaultModels) +} + +func (s *Service) catalogExpired(ctx context.Context, providerCfg config.ProviderConfig) bool { + catalog, err := s.loadCatalog(ctx, providerCfg) + if err != nil { + return false + } + return catalog.Expired(s.now()) +} + +func (s *Service) loadCatalogModels(ctx context.Context, providerCfg config.ProviderConfig) ([]ModelDescriptor, bool) { + catalog, err := s.loadCatalog(ctx, providerCfg) + if err != nil { + return nil, false + } + return catalog.Models, true +} + +func (s *Service) loadCatalog(ctx context.Context, providerCfg config.ProviderConfig) (ModelCatalog, error) { + if s.catalogs == nil { + return ModelCatalog{}, ErrModelCatalogNotFound + } + + identity, err := providerCfg.Identity() + if err != nil { + return ModelCatalog{}, err + } + return s.catalogs.Load(ctx, identity) +} + +func (s *Service) discoverAndPersist(ctx context.Context, providerCfg config.ProviderConfig) ([]ModelDescriptor, bool) { + if s.registry == nil { + return nil, false + } + if !s.registry.Supports(providerCfg.Driver) { + return nil, false + } + + resolved, err := providerCfg.Resolve() + if err != nil { + return nil, false + } + + discovered, err := s.registry.DiscoverModels(ctx, resolved) + if err != nil { + return nil, false + } + + discovered = MergeModelDescriptors(discovered) + if s.catalogs == nil { + return discovered, true + } + + identity, err := providerCfg.Identity() + if err != nil { + return discovered, true + } + + now := s.now() + _ = s.catalogs.Save(ctx, ModelCatalog{ + SchemaVersion: modelCatalogSchemaVersion, + Identity: identity, + FetchedAt: now, + ExpiresAt: now.Add(s.catalogTTL), + Models: discovered, + }) + return discovered, true +} + +func (s *Service) queueRefresh(providerCfg config.ProviderConfig) { + if s.catalogs == nil || s.registry == nil || !s.registry.Supports(providerCfg.Driver) { + return + } + + identity, err := providerCfg.Identity() + if err != nil { + return + } + + key := identity.Key() + s.refreshMu.Lock() + if _, exists := s.inFlightByID[key]; exists { + s.refreshMu.Unlock() + return + } + s.inFlightByID[key] = struct{}{} + s.refreshMu.Unlock() + + go func() { + defer func() { + s.refreshMu.Lock() + delete(s.inFlightByID, key) + s.refreshMu.Unlock() + }() + + ctx, cancel := context.WithTimeout(context.Background(), s.backgroundTimeout) + defer cancel() + _, _ = s.discoverAndPersist(ctx, providerCfg) + }() +} + func selectionFromConfig(cfg config.Config) ProviderSelection { return ProviderSelection{ ProviderID: cfg.SelectedProvider, @@ -139,48 +321,29 @@ func selectModel(currentModel string, models []string, fallback string) string { return strings.TrimSpace(fallback) } -func catalogItemFromConfig(cfg config.ProviderConfig) ProviderCatalogItem { +func catalogItemFromConfig(cfg config.ProviderConfig, models []ModelDescriptor) ProviderCatalogItem { return ProviderCatalogItem{ ID: strings.TrimSpace(cfg.Name), Name: strings.TrimSpace(cfg.Name), - Models: modelDescriptors(cfg.SupportedModels()), + Models: MergeModelDescriptors(models), } } -func modelDescriptors(models []string) []ModelDescriptor { +func modelDescriptorIDs(models []ModelDescriptor) []string { if len(models) == 0 { return nil } - descriptors := make([]ModelDescriptor, 0, len(models)) - for _, modelID := range models { - modelID = strings.TrimSpace(modelID) - if modelID == "" { + ids := make([]string, 0, len(models)) + for _, model := range models { + if strings.TrimSpace(model.ID) == "" { continue } - descriptors = append(descriptors, ModelDescriptor{ - ID: modelID, - Name: modelID, - }) + ids = append(ids, model.ID) } - return descriptors + return ids } -// Build 根据已解析的 provider 配置构建 Provider 实例。 -// 使 Service 满足 runtime.ProviderFactory 接口，简化装配层依赖。 -func (s *Service) Build(ctx context.Context, cfg config.ResolvedProviderConfig) (Provider, error) { - if err := s.validate(); err != nil { - return nil, err - } - return s.registry.Build(ctx, cfg) -} - -func (s *Service) validate() error { - if s == nil || s.manager == nil { - return errServiceManagerNil - } - if s.registry == nil { - return errServiceRegistryNil - } - return nil +func containsModelDescriptorID(models []ModelDescriptor, modelID string) bool { + return config.ContainsModelID(modelDescriptorIDs(models), modelID) } diff --git a/internal/provider/service_test.go b/internal/provider/service_test.go index 248dd31f..43e393b9 100644 --- a/internal/provider/service_test.go +++ b/internal/provider/service_test.go @@ -2,7 +2,9 @@ package provider import ( "context" + "sync" "testing" + "time" "neo-code/internal/config" ) @@ -10,337 +12,325 @@ import ( func TestNewService(t *testing.T) { t.Parallel() - t.Run("success", func(t *testing.T) { - t.Parallel() - manager := config.NewManager(config.NewLoader(t.TempDir(), testDefaultConfig())) - registry := NewRegistry() - service := NewService(manager, registry) - if service == nil { - t.Fatal("expected non-nil service") - } - }) - - t.Run("nil manager", func(t *testing.T) { - t.Parallel() - registry := NewRegistry() - service := NewService(nil, registry) - if service == nil { - t.Fatal("expected non-nil service even with nil manager") - } - }) + manager := config.NewManager(config.NewLoader(t.TempDir(), testDefaultConfig())) + registry := NewRegistry() + store := newMemoryCatalogStore() - t.Run("nil registry", func(t *testing.T) { - t.Parallel() - manager := config.NewManager(config.NewLoader(t.TempDir(), testDefaultConfig())) - service := NewService(manager, nil) - if service == nil { - t.Fatal("expected non-nil service even with nil registry") - } - }) + service := NewService(manager, registry, store) + if service == nil { + t.Fatal("expected non-nil service") + } + if service.catalogs != store { + t.Fatal("expected explicit catalog store to be used") + } } -func TestServiceListProviders(t *testing.T) { +func TestServiceListProvidersFiltersUnsupportedDrivers(t *testing.T) { t.Parallel() - t.Run("success", func(t *testing.T) { - t.Parallel() - manager := config.NewManager(config.NewLoader(t.TempDir(), testDefaultConfig())) - if _, err := manager.Load(context.Background()); err != nil { - t.Fatalf("Load() error = %v", err) - } + manager := config.NewManager(config.NewLoader(t.TempDir(), testDefaultConfig())) + if _, err := manager.Load(context.Background()); err != nil { + t.Fatalf("Load() error = %v", err) + } + if err := manager.Update(context.Background(), func(cfg *config.Config) error { + cfg.Providers = append(cfg.Providers, config.ProviderConfig{ + Name: "unsupported", + Driver: "custom", + BaseURL: "https://example.com/v1", + Model: "custom-model", + APIKeyEnv: "CUSTOM_API_KEY", + }) + return nil + }); err != nil { + t.Fatalf("append provider: %v", err) + } - registry := newServiceTestRegistry(t) - service := NewService(manager, registry) + registry := NewRegistry() + if err := registry.Register(testDriverDefinition(config.OpenAIName, nil)); err != nil { + t.Fatalf("register test driver: %v", err) + } - items, err := service.ListProviders(context.Background()) - if err != nil { - t.Fatalf("ListProviders() error = %v", err) - } - if len(items) == 0 { - t.Fatal("expected at least one provider") - } - if items[0].ID != testProviderName { - t.Fatalf("expected first provider ID %q, got %q", testProviderName, items[0].ID) + service := NewService(manager, registry, newMemoryCatalogStore()) + items, err := service.ListProviders(context.Background()) + if err != nil { + t.Fatalf("ListProviders() error = %v", err) + } + if len(items) != 4 { + t.Fatalf("expected 4 supported builtin providers, got %d", len(items)) + } + for _, item := range items { + if len(item.Models) != 1 { + t.Fatalf("expected default fallback model for %q, got %+v", item.ID, item.Models) } - }) + } +} - t.Run("nil service", func(t *testing.T) { - t.Parallel() - var service *Service - _, err := service.ListProviders(context.Background()) - if err == nil { - t.Fatal("expected error for nil service") - } - }) +func TestServiceSelectProviderFallsBackToProviderDefault(t *testing.T) { + t.Parallel() - t.Run("nil manager", func(t *testing.T) { - t.Parallel() - service := &Service{} - _, err := service.ListProviders(context.Background()) - if err == nil { - t.Fatal("expected error for nil manager") - } + defaults := testDefaultConfig() + defaults.Providers = append(defaults.Providers, config.ProviderConfig{ + Name: "custom-main", + Driver: "custom", + BaseURL: "https://example.com/v1", + Model: "custom-model", + Models: []string{"custom-model", "custom-alt"}, + APIKeyEnv: "CUSTOM_API_KEY", }) - t.Run("nil registry", func(t *testing.T) { - t.Parallel() - manager := config.NewManager(config.NewLoader(t.TempDir(), testDefaultConfig())) - if _, err := manager.Load(context.Background()); err != nil { - t.Fatalf("Load() error = %v", err) - } - service := &Service{manager: manager} - _, err := service.ListProviders(context.Background()) - if err == nil { - t.Fatal("expected error for nil registry") - } - }) + manager := config.NewManager(config.NewLoader(t.TempDir(), defaults)) + if _, err := manager.Load(context.Background()); err != nil { + t.Fatalf("Load() error = %v", err) + } + if err := manager.Update(context.Background(), func(cfg *config.Config) error { + cfg.CurrentModel = "gpt-5.4" + return nil + }); err != nil { + t.Fatalf("seed current model: %v", err) + } - t.Run("canceled context", func(t *testing.T) { - t.Parallel() - manager := config.NewManager(config.NewLoader(t.TempDir(), testDefaultConfig())) - if _, err := manager.Load(context.Background()); err != nil { - t.Fatalf("Load() error = %v", err) - } - registry := NewRegistry() - service := NewService(manager, registry) + registry := NewRegistry() + if err := registry.Register(testDriverDefinition("custom", nil)); err != nil { + t.Fatalf("register custom driver: %v", err) + } - ctx, cancel := context.WithCancel(context.Background()) - cancel() - _, err := service.ListProviders(ctx) - if err == nil { - t.Fatal("expected error for canceled context") - } - }) + service := NewService(manager, registry, newMemoryCatalogStore()) + selection, err := service.SelectProvider(context.Background(), "custom-main") + if err != nil { + t.Fatalf("SelectProvider() error = %v", err) + } + if selection.ProviderID != "custom-main" || selection.ModelID != "custom-model" { + t.Fatalf("unexpected selection: %+v", selection) + } } -func TestServiceSelectProvider(t *testing.T) { - t.Run("success", func(t *testing.T) { - t.Setenv(testAPIKeyEnv, "test-key") - - manager := config.NewManager(config.NewLoader(t.TempDir(), testDefaultConfig())) - if _, err := manager.Load(context.Background()); err != nil { - t.Fatalf("Load() error = %v", err) - } - - registry := newServiceTestRegistry(t) - service := NewService(manager, registry) - - selection, err := service.SelectProvider(context.Background(), testProviderName) - if err != nil { - t.Fatalf("SelectProvider() error = %v", err) - } - if selection.ProviderID != testProviderName { - t.Fatalf("expected provider ID %q, got %q", testProviderName, selection.ProviderID) - } - }) - - t.Run("provider not found", func(t *testing.T) { - t.Parallel() - manager := config.NewManager(config.NewLoader(t.TempDir(), testDefaultConfig())) - if _, err := manager.Load(context.Background()); err != nil { - t.Fatalf("Load() error = %v", err) - } - registry := NewRegistry() - service := NewService(manager, registry) +func TestServiceListModelsUsesConfiguredModelsWithoutDiscovery(t *testing.T) { + t.Parallel() - _, err := service.SelectProvider(context.Background(), "nonexistent") - if err == nil { - t.Fatal("expected error for nonexistent provider") - } - }) + manager := config.NewManager(config.NewLoader(t.TempDir(), testDefaultConfig())) + if _, err := manager.Load(context.Background()); err != nil { + t.Fatalf("Load() error = %v", err) + } + if err := manager.Update(context.Background(), func(cfg *config.Config) error { + cfg.Providers = append(cfg.Providers, config.ProviderConfig{ + Name: "broken-provider", + Driver: "missing-driver", + BaseURL: "https://example.com/v1", + Model: "broken-model", + Models: []string{"broken-model", "broken-alt"}, + APIKeyEnv: "BROKEN_API_KEY", + }) + cfg.SelectedProvider = "broken-provider" + cfg.CurrentModel = "broken-model" + return nil + }); err != nil { + t.Fatalf("append provider: %v", err) + } - t.Run("nil service", func(t *testing.T) { - t.Parallel() - var service *Service - _, err := service.SelectProvider(context.Background(), testProviderName) - if err == nil { - t.Fatal("expected error for nil service") - } - }) + service := NewService(manager, NewRegistry(), newMemoryCatalogStore()) + models, err := service.ListModels(context.Background()) + if err != nil { + t.Fatalf("ListModels() error = %v", err) + } + if len(models) != 2 || models[1].ID != "broken-alt" { + t.Fatalf("expected configured models fallback, got %+v", models) + } } -func TestServiceListModels(t *testing.T) { - t.Run("success", func(t *testing.T) { - t.Setenv(testAPIKeyEnv, "test-key") +func TestServiceSetCurrentModelUsesConfiguredModels(t *testing.T) { + t.Parallel() - manager := config.NewManager(config.NewLoader(t.TempDir(), testDefaultConfig())) - if _, err := manager.Load(context.Background()); err != nil { - t.Fatalf("Load() error = %v", err) - } + manager := config.NewManager(config.NewLoader(t.TempDir(), testDefaultConfig())) + if _, err := manager.Load(context.Background()); err != nil { + t.Fatalf("Load() error = %v", err) + } + if err := manager.Update(context.Background(), func(cfg *config.Config) error { + cfg.Providers[0].Models = []string{cfg.Providers[0].Model, "gpt-4o"} + return nil + }); err != nil { + t.Fatalf("seed provider models: %v", err) + } - registry := NewRegistry() - service := NewService(manager, registry) + service := NewService(manager, NewRegistry(), newMemoryCatalogStore()) + selection, err := service.SetCurrentModel(context.Background(), "gpt-4o") + if err != nil { + t.Fatalf("SetCurrentModel() error = %v", err) + } + if selection.ModelID != "gpt-4o" { + t.Fatalf("expected selected model %q, got %+v", "gpt-4o", selection) + } +} - models, err := service.ListModels(context.Background()) - if err != nil { - t.Fatalf("ListModels() error = %v", err) - } - if len(models) == 0 { - t.Fatal("expected at least one model") - } - }) +func TestServiceListModelsDiscoversAndCachesOnMiss(t *testing.T) { + t.Setenv(testAPIKeyEnv, "test-key") - t.Run("nil service", func(t *testing.T) { - t.Parallel() - var service *Service - _, err := service.ListModels(context.Background()) - if err == nil { - t.Fatal("expected error for nil service") - } - }) + manager := config.NewManager(config.NewLoader(t.TempDir(), testDefaultConfig())) + if _, err := manager.Load(context.Background()); err != nil { + t.Fatalf("Load() error = %v", err) + } - t.Run("canceled context", func(t *testing.T) { - t.Setenv(testAPIKeyEnv, "test-key") + store := newMemoryCatalogStore() + registry := NewRegistry() + if err := registry.Register(testDriverDefinition(config.OpenAIName, func(ctx context.Context, cfg config.ResolvedProviderConfig) ([]ModelDescriptor, error) { + return []ModelDescriptor{{ + ID: "server-model", + Name: "Server Model", + ContextWindow: 32000, + MaxOutputTokens: 4096, + Metadata: map[string]any{ + "id": "server-model", + "context_window": float64(32000), + "max_output_tokens": float64(4096), + }, + }}, nil + })); err != nil { + t.Fatalf("register discovery driver: %v", err) + } - manager := config.NewManager(config.NewLoader(t.TempDir(), testDefaultConfig())) - if _, err := manager.Load(context.Background()); err != nil { - t.Fatalf("Load() error = %v", err) - } - registry := NewRegistry() - service := NewService(manager, registry) + service := NewService(manager, registry, store) + models, err := service.ListModels(context.Background()) + if err != nil { + t.Fatalf("ListModels() error = %v", err) + } + if !containsModelDescriptorID(models, "server-model") { + t.Fatalf("expected discovered model in result, got %+v", models) + } - ctx, cancel := context.WithCancel(context.Background()) - cancel() - _, err := service.ListModels(ctx) - if err == nil { - t.Fatal("expected error for canceled context") - } - }) + identity, err := config.OpenAIProvider().Identity() + if err != nil { + t.Fatalf("Identity() error = %v", err) + } + catalog, err := store.Load(context.Background(), identity) + if err != nil { + t.Fatalf("Load() cached catalog error = %v", err) + } + if !containsModelDescriptorID(catalog.Models, "server-model") { + t.Fatalf("expected cached discovered model, got %+v", catalog.Models) + } } -func TestServiceSetCurrentModel(t *testing.T) { - t.Run("success", func(t *testing.T) { - t.Setenv(testAPIKeyEnv, "test-key") +func TestServiceListModelsReturnsStaleCacheAndRefreshesInBackground(t *testing.T) { + t.Setenv("CUSTOM_API_KEY", "test-key") + + defaults := testDefaultConfig() + defaults.Providers = []config.ProviderConfig{{ + Name: "custom", + Driver: "custom", + BaseURL: "https://example.com/v1", + Model: "fallback-model", + APIKeyEnv: "CUSTOM_API_KEY", + }} + defaults.SelectedProvider = "custom" + defaults.CurrentModel = "fallback-model" + + manager := config.NewManager(config.NewLoader(t.TempDir(), defaults)) + if _, err := manager.Load(context.Background()); err != nil { + t.Fatalf("Load() error = %v", err) + } - manager := config.NewManager(config.NewLoader(t.TempDir(), testDefaultConfig())) - if _, err := manager.Load(context.Background()); err != nil { - t.Fatalf("Load() error = %v", err) - } + identity, err := defaults.Providers[0].Identity() + if err != nil { + t.Fatalf("Identity() error = %v", err) + } - registry := NewRegistry() - service := NewService(manager, registry) + store := newMemoryCatalogStore() + now := time.Date(2026, 4, 2, 12, 0, 0, 0, time.UTC) + if err := store.Save(context.Background(), ModelCatalog{ + SchemaVersion: modelCatalogSchemaVersion, + Identity: identity, + FetchedAt: now.Add(-48 * time.Hour), + ExpiresAt: now.Add(-24 * time.Hour), + Models: []ModelDescriptor{{ + ID: "stale-model", + Name: "Stale Model", + }}, + }); err != nil { + t.Fatalf("seed stale catalog: %v", err) + } - newModel := "gpt-4o" - selection, err := service.SetCurrentModel(context.Background(), newModel) - if err != nil { - t.Fatalf("SetCurrentModel() error = %v", err) - } - if selection.ModelID != newModel { - t.Fatalf("expected model ID %q, got %q", newModel, selection.ModelID) - } - }) + refreshed := make(chan struct{}, 1) + registry := NewRegistry() + if err := registry.Register(testDriverDefinition("custom", func(ctx context.Context, cfg config.ResolvedProviderConfig) ([]ModelDescriptor, error) { + select { + case refreshed <- struct{}{}: + default: + } + return []ModelDescriptor{{ID: "fresh-model", Name: "Fresh Model"}}, nil + })); err != nil { + t.Fatalf("register custom driver: %v", err) + } - t.Run("empty model ID", func(t *testing.T) { - t.Parallel() - manager := config.NewManager(config.NewLoader(t.TempDir(), testDefaultConfig())) - if _, err := manager.Load(context.Background()); err != nil { - t.Fatalf("Load() error = %v", err) - } - registry := NewRegistry() - service := NewService(manager, registry) + service := NewService(manager, registry, store) + service.now = func() time.Time { return now } + service.backgroundTimeout = time.Second - _, err := service.SetCurrentModel(context.Background(), "") - if err == nil { - t.Fatal("expected error for empty model ID") - } - }) + models, err := service.ListModels(context.Background()) + if err != nil { + t.Fatalf("ListModels() error = %v", err) + } + if !containsModelDescriptorID(models, "stale-model") { + t.Fatalf("expected stale cached model to be returned immediately, got %+v", models) + } - t.Run("model not found", func(t *testing.T) { - t.Parallel() - manager := config.NewManager(config.NewLoader(t.TempDir(), testDefaultConfig())) - if _, err := manager.Load(context.Background()); err != nil { - t.Fatalf("Load() error = %v", err) - } - registry := NewRegistry() - service := NewService(manager, registry) + select { + case <-refreshed: + case <-time.After(2 * time.Second): + t.Fatal("expected background refresh to run") + } - _, err := service.SetCurrentModel(context.Background(), "nonexistent-model") - if err == nil { - t.Fatal("expected error for nonexistent model") + deadline := time.Now().Add(2 * time.Second) + for time.Now().Before(deadline) { + catalog, err := store.Load(context.Background(), identity) + if err == nil && containsModelDescriptorID(catalog.Models, "fresh-model") { + return } - }) + time.Sleep(20 * time.Millisecond) + } - t.Run("nil service", func(t *testing.T) { - t.Parallel() - var service *Service - _, err := service.SetCurrentModel(context.Background(), "model") - if err == nil { - t.Fatal("expected error for nil service") - } - }) + catalog, err := store.Load(context.Background(), identity) + if err != nil { + t.Fatalf("Load() refreshed catalog error = %v", err) + } + t.Fatalf("expected refreshed catalog to contain fresh-model, got %+v", catalog.Models) } -func TestServiceBuild(t *testing.T) { +func TestServiceBuildAndValidate(t *testing.T) { t.Parallel() - t.Run("nil service", func(t *testing.T) { + t.Run("build delegates to registry", func(t *testing.T) { t.Parallel() - var service *Service - _, err := service.Build(context.Background(), config.ResolvedProviderConfig{}) - if err == nil { - t.Fatal("expected error for nil service") - } - }) - - t.Run("nil manager", func(t *testing.T) { - t.Parallel() - service := &Service{} - _, err := service.Build(context.Background(), config.ResolvedProviderConfig{}) - if err == nil { - t.Fatal("expected error for nil manager") - } - }) - t.Run("nil registry", func(t *testing.T) { - t.Parallel() manager := config.NewManager(config.NewLoader(t.TempDir(), testDefaultConfig())) - service := &Service{manager: manager} - _, err := service.Build(context.Background(), config.ResolvedProviderConfig{}) - if err == nil { - t.Fatal("expected error for nil registry") + registry := NewRegistry() + if err := registry.Register(testDriverDefinition("custom", nil)); err != nil { + t.Fatalf("register driver: %v", err) + } + + service := NewService(manager, registry, newMemoryCatalogStore()) + providerInstance, err := service.Build(context.Background(), config.ResolvedProviderConfig{ + ProviderConfig: config.ProviderConfig{ + Name: "custom", + Driver: "custom", + BaseURL: "https://example.com/v1", + Model: "model", + APIKeyEnv: "CUSTOM_API_KEY", + }, + APIKey: "test-key", + }) + if err != nil { + t.Fatalf("Build() error = %v", err) } - }) -} - -func TestServiceValidate(t *testing.T) { - t.Parallel() - - t.Run("nil service", func(t *testing.T) { - t.Parallel() - var service *Service - err := service.validate() - if err == nil { - t.Fatal("expected error for nil service") + if _, ok := providerInstance.(serviceTestProvider); !ok { + t.Fatalf("expected serviceTestProvider, got %T", providerInstance) } }) - t.Run("nil manager", func(t *testing.T) { + t.Run("nil service fails validate", func(t *testing.T) { t.Parallel() - service := &Service{} - err := service.validate() - if err == nil { - t.Fatal("expected error for nil manager") - } - }) - t.Run("nil registry", func(t *testing.T) { - t.Parallel() - manager := config.NewManager(config.NewLoader(t.TempDir(), testDefaultConfig())) - service := &Service{manager: manager} - err := service.validate() - if err == nil { - t.Fatal("expected error for nil registry") - } - }) - - t.Run("success", func(t *testing.T) { - t.Parallel() - manager := config.NewManager(config.NewLoader(t.TempDir(), testDefaultConfig())) - registry := NewRegistry() - service := NewService(manager, registry) - err := service.validate() - if err != nil { - t.Fatalf("expected no error, got %v", err) + var service *Service + if err := service.validate(); err == nil { + t.Fatal("expected validate error for nil service") } }) } @@ -369,69 +359,28 @@ func TestSelectModelHelper(t *testing.T) { fallback: "gpt-4.1", expected: "gpt-4.1", }, - { - name: "empty current model", - currentModel: "", - models: []string{"gpt-4.1", "gpt-4o", "gpt-5.4"}, - fallback: "gpt-4.1", - expected: "gpt-4.1", - }, - { - name: "current model with whitespace", - currentModel: " gpt-4o ", - models: []string{"gpt-4.1", "gpt-4o", "gpt-5.4"}, - fallback: "gpt-4.1", - expected: "gpt-4o", - }, } for _, tt := range tests { t.Run(tt.name, func(t *testing.T) { t.Parallel() - result := selectModel(tt.currentModel, tt.models, tt.fallback) - if result != tt.expected { - t.Fatalf("expected %q, got %q", tt.expected, result) + if got := selectModel(tt.currentModel, tt.models, tt.fallback); got != tt.expected { + t.Fatalf("selectModel() = %q, want %q", got, tt.expected) } }) } } -func TestModelDescriptorsHelper(t *testing.T) { +func TestModelDescriptorsFromIDsHelper(t *testing.T) { t.Parallel() - t.Run("empty models", func(t *testing.T) { - t.Parallel() - result := modelDescriptors(nil) - if result != nil { - t.Fatalf("expected nil for empty models, got %v", result) - } - }) - - t.Run("non-empty models", func(t *testing.T) { - t.Parallel() - models := []string{"gpt-4.1", "gpt-4o", "gpt-5.4"} - result := modelDescriptors(models) - if len(result) != 3 { - t.Fatalf("expected 3 descriptors, got %d", len(result)) - } - for i, desc := range result { - if desc.ID != models[i] { - t.Fatalf("expected ID %q, got %q", models[i], desc.ID) - } - if desc.Name != models[i] { - t.Fatalf("expected name %q, got %q", models[i], desc.Name) - } - } - }) - - t.Run("models with empty strings", func(t *testing.T) { - t.Parallel() - models := []string{"gpt-4.1", "", "gpt-4o", " ", "gpt-5.4"} - result := modelDescriptors(models) - if len(result) != 3 { - t.Fatalf("expected 3 descriptors after filtering empty strings, got %d", len(result)) - } - }) + models := modelDescriptorsFromIDs([]string{"gpt-4.1", "", "gpt-4o"}) + if len(models) != 2 { + t.Fatalf("expected 2 descriptors, got %d", len(models)) + } + if models[0].ID != "gpt-4.1" || models[1].ID != "gpt-4o" { + t.Fatalf("unexpected descriptors: %+v", models) + } } func testDefaultConfig() *config.Config { @@ -445,19 +394,14 @@ func testDefaultConfig() *config.Config { return cfg } -func newServiceTestRegistry(t *testing.T) *Registry { - t.Helper() - - registry := NewRegistry() - if err := registry.Register(DriverDefinition{ - Name: testProviderName, +func testDriverDefinition(name string, discover DiscoveryFunc) DriverDefinition { + return DriverDefinition{ + Name: name, + Discover: discover, Build: func(ctx context.Context, cfg config.ResolvedProviderConfig) (Provider, error) { return serviceTestProvider{}, nil }, - }); err != nil { - t.Fatalf("register test driver: %v", err) } - return registry } type serviceTestProvider struct{} @@ -466,6 +410,44 @@ func (serviceTestProvider) Chat(ctx context.Context, req ChatRequest, events cha return ChatResponse{}, nil } +type memoryCatalogStore struct { + mu sync.Mutex + catalogs map[string]ModelCatalog +} + +func newMemoryCatalogStore() *memoryCatalogStore { + return &memoryCatalogStore{ + catalogs: map[string]ModelCatalog{}, + } +} + +func (s *memoryCatalogStore) Load(ctx context.Context, identity config.ProviderIdentity) (ModelCatalog, error) { + if err := ctx.Err(); err != nil { + return ModelCatalog{}, err + } + + s.mu.Lock() + defer s.mu.Unlock() + + catalog, ok := s.catalogs[identity.Key()] + if !ok { + return ModelCatalog{}, ErrModelCatalogNotFound + } + return catalog, nil +} + +func (s *memoryCatalogStore) Save(ctx context.Context, catalog ModelCatalog) error { + if err := ctx.Err(); err != nil { + return err + } + + s.mu.Lock() + defer s.mu.Unlock() + + s.catalogs[catalog.Identity.Key()] = catalog + return nil +} + const ( testProviderName = "openai" testAPIKeyEnv = "OPENAI_API_KEY" diff --git a/internal/provider/types.go b/internal/provider/types.go index 2a08b2f0..a46d6fe7 100644 --- a/internal/provider/types.go +++ b/internal/provider/types.go @@ -1,9 +1,5 @@ package provider -import ( - "errors" -) - const ( RoleSystem = "system" RoleUser = "user" @@ -50,16 +46,14 @@ type Usage struct { TotalTokens int `json:"total_tokens"` } -var ( - ErrProviderNotFound = errors.New("provider not found") - ErrModelNotFound = errors.New("model not found") - ErrDriverNotFound = errors.New("provider driver not found") -) - type ModelDescriptor struct { - ID string `json:"id"` - Name string `json:"name"` - Description string `json:"description,omitempty"` + ID string `json:"id"` + Name string `json:"name"` + Description string `json:"description,omitempty"` + ContextWindow int `json:"context_window,omitempty"` + MaxOutputTokens int `json:"max_output_tokens,omitempty"` + Capabilities map[string]bool `json:"capabilities,omitempty"` + Metadata map[string]any `json:"metadata,omitempty"` } type ProviderCatalogItem struct { diff --git a/internal/tui/update_test.go b/internal/tui/update_test.go index 5220507a..a1e649ed 100644 --- a/internal/tui/update_test.go +++ b/internal/tui/update_test.go @@ -928,6 +928,7 @@ func TestAppUpdateAdditionalTransitions(t *testing.T) { func TestAppUpdateModelPickerEnterAppliesSelection(t *testing.T) { manager := newTestConfigManager(t) if err := manager.Update(context.Background(), func(cfg *config.Config) error { + cfg.Providers[0].Models = []string{cfg.Providers[0].Model, "gpt-4o"} cfg.CurrentModel = "unsupported-current" return nil }); err != nil { @@ -952,10 +953,12 @@ func TestAppUpdateModelPickerEnterAppliesSelection(t *testing.T) { t.Fatalf("expected picker to close after selection") } - for _, msg := range collectTeaMessages(cmd) { - model, follow := app.Update(msg) - app = model.(App) - _ = collectTeaMessages(follow) + if cmd != nil { + if msg := cmd(); msg != nil { + model, follow := app.Update(msg) + app = model.(App) + _ = collectTeaMessages(follow) + } } cfg := manager.Get() @@ -970,7 +973,7 @@ func TestAppUpdateProviderPickerEnterAppliesSelection(t *testing.T) { cfg.Providers = append(cfg.Providers, config.ProviderConfig{ Name: "openai-alt", Driver: "openai", - BaseURL: config.OpenAIDefaultBaseURL, + BaseURL: "https://alt.example.com/v1", Model: "gpt-4o", Models: []string{"gpt-4o"}, APIKeyEnv: config.OpenAIDefaultAPIKeyEnv, @@ -1489,7 +1492,7 @@ func newTestProviderService(t *testing.T, manager *config.Manager) *provider.Ser if err != nil { t.Fatalf("register provider drivers: %v", err) } - return provider.NewService(manager, registry) + return provider.NewService(manager, registry, nil) } func collectTeaMessages(cmd tea.Cmd) []tea.Msg { @@ -1504,7 +1507,7 @@ func collectTeaMessages(cmd tea.Cmd) []tea.Msg { var msg tea.Msg select { case msg = <-msgCh: - case <-time.After(25 * time.Millisecond): + case <-time.After(250 * time.Millisecond): return nil } if msg == nil { From a70f43d4c3e81a973e095129f6c5d4c4287fa6b1 Mon Sep 17 00:00:00 2001 From: Cai_Tang <106404101+Cai-Tang-www@users.noreply.github.com> Date: Thu, 2 Apr 2026 09:20:20 +0800 Subject: [PATCH 07/23] =?UTF-8?q?app=EF=BC=9A=E6=96=B0=E5=A2=9E=20workspac?= =?UTF-8?q?e=20sandbox=20=E5=B7=A5=E5=85=B7=E6=9D=83=E9=99=90=E6=8E=A7?= =?UTF-8?q?=E5=88=B6?= MIME-Version: 1.0 Content-Type: text/plain; charset=UTF-8 Content-Transfer-Encoding: 8bit --- internal/security/types.go | 16 +- internal/security/workspace.go | 220 ++++++++++++++++++++++++++++ internal/security/workspace_test.go | 217 +++++++++++++++++++++++++++ internal/tools/manager.go | 6 + internal/tools/manager_test.go | 49 ++++++- internal/tools/permission_mapper.go | 12 ++ 6 files changed, 512 insertions(+), 8 deletions(-) create mode 100644 internal/security/workspace.go create mode 100644 internal/security/workspace_test.go diff --git a/internal/security/types.go b/internal/security/types.go index 95a32e6a..820d74fe 100644 --- a/internal/security/types.go +++ b/internal/security/types.go @@ -58,13 +58,15 @@ const ( // ActionPayload is the normalized structured context used by policy and sandbox. type ActionPayload struct { - ToolName string - Resource string - Operation string - SessionID string - Workdir string - TargetType TargetType - Target string + ToolName string + Resource string + Operation string + SessionID string + Workdir string + TargetType TargetType + Target string + SandboxTargetType TargetType + SandboxTarget string } // Action is the unified security input for one tool execution request. diff --git a/internal/security/workspace.go b/internal/security/workspace.go new file mode 100644 index 00000000..65335f35 --- /dev/null +++ b/internal/security/workspace.go @@ -0,0 +1,220 @@ +package security + +import ( + "context" + "errors" + "fmt" + "os" + "path/filepath" + "strings" +) + +// WorkspaceSandbox enforces workspace-relative path boundaries for tool actions. +type WorkspaceSandbox struct{} + +// NewWorkspaceSandbox creates a sandbox that blocks traversal and symlink escape. +func NewWorkspaceSandbox() *WorkspaceSandbox { + return &WorkspaceSandbox{} +} + +// Check validates that the action stays within the configured workspace root. +func (s *WorkspaceSandbox) Check(ctx context.Context, action Action) error { + if err := ctx.Err(); err != nil { + return err + } + if err := action.Validate(); err != nil { + return err + } + + plan, ok, err := buildWorkspacePlan(action) + if err != nil { + return err + } + if !ok { + return nil + } + + return validateWorkspacePlan(plan) +} + +type workspacePlan struct { + root string + target string +} + +func buildWorkspacePlan(action Action) (workspacePlan, bool, error) { + if !needsWorkspaceSandbox(action) { + return workspacePlan{}, false, nil + } + + root := strings.TrimSpace(action.Payload.Workdir) + if root == "" { + return workspacePlan{}, false, errors.New("security: workspace root is empty") + } + + target, ok := sandboxTarget(action) + if !ok { + return workspacePlan{}, false, nil + } + + return workspacePlan{ + root: root, + target: target, + }, true, nil +} + +func needsWorkspaceSandbox(action Action) bool { + switch action.Type { + case ActionTypeRead, ActionTypeWrite, ActionTypeBash: + return true + default: + return false + } +} + +func sandboxTarget(action Action) (string, bool) { + if action.Type == ActionTypeBash { + target := strings.TrimSpace(action.Payload.SandboxTarget) + if target == "" { + return ".", true + } + return target, true + } + + targetType := action.Payload.SandboxTargetType + if targetType == "" { + targetType = action.Payload.TargetType + } + + target := strings.TrimSpace(action.Payload.SandboxTarget) + if target == "" { + target = strings.TrimSpace(action.Payload.Target) + } + + switch targetType { + case TargetTypeDirectory: + if target == "" { + return ".", true + } + return target, true + case TargetTypePath: + if target == "" { + return "", false + } + return target, true + default: + return "", false + } +} + +func validateWorkspacePlan(plan workspacePlan) error { + root, err := canonicalWorkspaceRoot(plan.root) + if err != nil { + return err + } + + target, err := absoluteWorkspaceTarget(root, plan.target) + if err != nil { + return err + } + if !isWithinWorkspace(root, target) { + return fmt.Errorf("security: path %q escapes workspace root", plan.target) + } + + return ensureNoSymlinkEscape(root, target, plan.target) +} + +func canonicalWorkspaceRoot(root string) (string, error) { + absoluteRoot, err := filepath.Abs(strings.TrimSpace(root)) + if err != nil { + return "", fmt.Errorf("security: resolve workspace root: %w", err) + } + + canonicalRoot, err := filepath.EvalSymlinks(absoluteRoot) + if err != nil { + if errors.Is(err, os.ErrNotExist) { + return filepath.Clean(absoluteRoot), nil + } + return "", fmt.Errorf("security: resolve workspace root: %w", err) + } + + return filepath.Clean(canonicalRoot), nil +} + +func absoluteWorkspaceTarget(root string, target string) (string, error) { + trimmedTarget := strings.TrimSpace(target) + if trimmedTarget == "" { + trimmedTarget = "." + } + if !filepath.IsAbs(trimmedTarget) { + trimmedTarget = filepath.Join(root, trimmedTarget) + } + + absoluteTarget, err := filepath.Abs(trimmedTarget) + if err != nil { + return "", fmt.Errorf("security: resolve workspace target %q: %w", target, err) + } + + return filepath.Clean(absoluteTarget), nil +} + +func ensureNoSymlinkEscape(root string, target string, original string) error { + relativeTarget, err := filepath.Rel(root, target) + if err != nil { + return fmt.Errorf("security: compare workspace target %q: %w", original, err) + } + + cleanRelative := filepath.Clean(relativeTarget) + if cleanRelative == "." { + return nil + } + + current := root + for _, segment := range splitRelativePath(cleanRelative) { + next := filepath.Join(current, segment) + info, err := os.Lstat(next) + if err != nil { + if errors.Is(err, os.ErrNotExist) { + return nil + } + return fmt.Errorf("security: inspect path %q: %w", next, err) + } + + if info.Mode()&os.ModeSymlink != 0 { + resolved, err := filepath.EvalSymlinks(next) + if err != nil { + return fmt.Errorf("security: resolve symlink %q: %w", next, err) + } + resolved, err = filepath.Abs(resolved) + if err != nil { + return fmt.Errorf("security: resolve symlink %q: %w", next, err) + } + if !isWithinWorkspace(root, resolved) { + return fmt.Errorf("security: path %q escapes workspace root via symlink", original) + } + current = filepath.Clean(resolved) + continue + } + + current = next + } + + return nil +} + +func splitRelativePath(path string) []string { + cleanPath := filepath.Clean(path) + if cleanPath == "." { + return nil + } + return strings.Split(cleanPath, string(os.PathSeparator)) +} + +func isWithinWorkspace(root string, target string) bool { + relativePath, err := filepath.Rel(root, target) + if err != nil { + return false + } + return relativePath == "." || + (relativePath != ".." && !strings.HasPrefix(relativePath, ".."+string(os.PathSeparator))) +} diff --git a/internal/security/workspace_test.go b/internal/security/workspace_test.go new file mode 100644 index 00000000..2a811826 --- /dev/null +++ b/internal/security/workspace_test.go @@ -0,0 +1,217 @@ +package security + +import ( + "context" + "os" + "path/filepath" + "strings" + "testing" +) + +func TestWorkspaceSandboxCheck(t *testing.T) { + t.Parallel() + + tests := []struct { + name string + prepare func(t *testing.T, root string, outside string) + action func(root string, outside string) Action + expectErr string + }{ + { + name: "read path inside workspace is allowed", + prepare: func(t *testing.T, root string, outside string) { + t.Helper() + mustWriteWorkspaceFile(t, filepath.Join(root, "notes.txt"), "hello") + }, + action: func(root string, outside string) Action { + return fileAction(ActionTypeRead, "filesystem_read_file", "read_file", root, "notes.txt") + }, + }, + { + name: "read traversal is rejected", + action: func(root string, outside string) Action { + return fileAction(ActionTypeRead, "filesystem_read_file", "read_file", root, filepath.Join("..", "outside.txt")) + }, + expectErr: "escapes workspace root", + }, + { + name: "absolute path outside workspace is rejected", + action: func(root string, outside string) Action { + return fileAction(ActionTypeRead, "filesystem_read_file", "read_file", root, outside) + }, + expectErr: "escapes workspace root", + }, + { + name: "symlinked file outside workspace is rejected", + prepare: func(t *testing.T, root string, outside string) { + t.Helper() + mustWriteWorkspaceFile(t, outside, "secret") + mustSymlinkOrSkip(t, outside, filepath.Join(root, "linked.txt")) + }, + action: func(root string, outside string) Action { + return fileAction(ActionTypeRead, "filesystem_read_file", "read_file", root, "linked.txt") + }, + expectErr: "via symlink", + }, + { + name: "symlinked parent directory outside workspace is rejected", + prepare: func(t *testing.T, root string, outside string) { + t.Helper() + outsideDir := filepath.Dir(outside) + if err := os.MkdirAll(outsideDir, 0o755); err != nil { + t.Fatalf("mkdir outside dir: %v", err) + } + mustSymlinkOrSkip(t, outsideDir, filepath.Join(root, "linked-dir")) + }, + action: func(root string, outside string) Action { + return fileAction(ActionTypeWrite, "filesystem_write_file", "write_file", root, filepath.Join("linked-dir", "new.txt")) + }, + expectErr: "via symlink", + }, + { + name: "missing nested write path inside workspace is allowed", + action: func(root string, outside string) Action { + return fileAction(ActionTypeWrite, "filesystem_write_file", "write_file", root, filepath.Join("new", "nested.txt")) + }, + }, + { + name: "grep defaults to workspace root when dir is empty", + action: func(root string, outside string) Action { + return Action{ + Type: ActionTypeRead, + Payload: ActionPayload{ + ToolName: "filesystem_grep", + Resource: "filesystem_grep", + Operation: "grep", + Workdir: root, + TargetType: TargetTypeDirectory, + SandboxTargetType: TargetTypeDirectory, + }, + } + }, + }, + { + name: "bash workdir inside workspace is allowed", + prepare: func(t *testing.T, root string, outside string) { + t.Helper() + if err := os.MkdirAll(filepath.Join(root, "scripts"), 0o755); err != nil { + t.Fatalf("mkdir scripts: %v", err) + } + }, + action: func(root string, outside string) Action { + return bashAction(root, "pwd", "scripts") + }, + }, + { + name: "bash workdir traversal is rejected", + action: func(root string, outside string) Action { + return bashAction(root, "pwd", filepath.Join("..", "outside")) + }, + expectErr: "escapes workspace root", + }, + { + name: "webfetch does not trigger workspace checks", + action: func(root string, outside string) Action { + return Action{ + Type: ActionTypeRead, + Payload: ActionPayload{ + ToolName: "webfetch", + Resource: "webfetch", + Operation: "fetch", + Workdir: root, + TargetType: TargetTypeURL, + Target: "https://example.com", + }, + } + }, + }, + { + name: "missing workspace root is rejected for path action", + action: func(root string, outside string) Action { + return fileAction(ActionTypeRead, "filesystem_read_file", "read_file", "", "notes.txt") + }, + expectErr: "workspace root is empty", + }, + { + name: "empty file target is deferred to tool validation", + action: func(root string, outside string) Action { + return fileAction(ActionTypeRead, "filesystem_read_file", "read_file", root, "") + }, + }, + } + + for _, tt := range tests { + tt := tt + t.Run(tt.name, func(t *testing.T) { + t.Parallel() + + root := t.TempDir() + outsideRoot := t.TempDir() + outsideFile := filepath.Join(outsideRoot, "outside.txt") + if tt.prepare != nil { + tt.prepare(t, root, outsideFile) + } + + sandbox := NewWorkspaceSandbox() + err := sandbox.Check(context.Background(), tt.action(root, outsideFile)) + if tt.expectErr != "" { + if err == nil || !strings.Contains(err.Error(), tt.expectErr) { + t.Fatalf("expected error containing %q, got %v", tt.expectErr, err) + } + return + } + if err != nil { + t.Fatalf("unexpected error: %v", err) + } + }) + } +} + +func fileAction(actionType ActionType, toolName string, operation string, workdir string, target string) Action { + return Action{ + Type: actionType, + Payload: ActionPayload{ + ToolName: toolName, + Resource: toolName, + Operation: operation, + Workdir: workdir, + TargetType: TargetTypePath, + Target: target, + SandboxTargetType: TargetTypePath, + SandboxTarget: target, + }, + } +} + +func bashAction(workdir string, command string, requestedWorkdir string) Action { + return Action{ + Type: ActionTypeBash, + Payload: ActionPayload{ + ToolName: "bash", + Resource: "bash", + Operation: "command", + Workdir: workdir, + TargetType: TargetTypeCommand, + Target: command, + SandboxTargetType: TargetTypeDirectory, + SandboxTarget: requestedWorkdir, + }, + } +} + +func mustWriteWorkspaceFile(t *testing.T, path string, content string) { + t.Helper() + if err := os.MkdirAll(filepath.Dir(path), 0o755); err != nil { + t.Fatalf("mkdir %s: %v", filepath.Dir(path), err) + } + if err := os.WriteFile(path, []byte(content), 0o644); err != nil { + t.Fatalf("write %s: %v", path, err) + } +} + +func mustSymlinkOrSkip(t *testing.T, target string, link string) { + t.Helper() + if err := os.Symlink(target, link); err != nil { + t.Skipf("symlink not supported in this environment: %v", err) + } +} diff --git a/internal/tools/manager.go b/internal/tools/manager.go index d7535f1e..fea81eb7 100644 --- a/internal/tools/manager.go +++ b/internal/tools/manager.go @@ -208,6 +208,12 @@ func actionMetadata(action security.Action) map[string]any { if action.Payload.Target != "" { metadata["permission_target"] = action.Payload.Target } + if action.Payload.SandboxTargetType != "" { + metadata["permission_sandbox_target_type"] = string(action.Payload.SandboxTargetType) + } + if action.Payload.SandboxTarget != "" { + metadata["permission_sandbox_target"] = action.Payload.SandboxTarget + } return metadata } diff --git a/internal/tools/manager_test.go b/internal/tools/manager_test.go index 7dcda95c..9f2615ec 100644 --- a/internal/tools/manager_test.go +++ b/internal/tools/manager_test.go @@ -3,6 +3,8 @@ package tools import ( "context" "errors" + "os" + "path/filepath" "strings" "testing" @@ -304,6 +306,41 @@ func TestDefaultManagerExecuteBoundaries(t *testing.T) { } } +func TestDefaultManagerExecuteWithWorkspaceSandbox(t *testing.T) { + t.Parallel() + + registry := NewRegistry() + tool := &managerStubTool{name: "filesystem_write_file", content: "ok"} + registry.Register(tool) + + engine, err := security.NewStaticGateway(security.DecisionAllow, nil) + if err != nil { + t.Fatalf("new engine: %v", err) + } + manager, err := NewManager(registry, engine, security.NewWorkspaceSandbox()) + if err != nil { + t.Fatalf("new manager: %v", err) + } + + workdir := t.TempDir() + outsideDir := t.TempDir() + if err := os.Symlink(outsideDir, filepath.Join(workdir, "link")); err != nil { + t.Skipf("symlink not supported in this environment: %v", err) + } + + _, execErr := manager.Execute(context.Background(), ToolCallInput{ + Name: "filesystem_write_file", + Arguments: []byte(`{"path":"link/outside.txt","content":"hello"}`), + Workdir: workdir, + }) + if execErr == nil || !strings.Contains(execErr.Error(), "escapes workspace root via symlink") { + t.Fatalf("expected sandbox escape error, got %v", execErr) + } + if tool.callCount != 0 { + t.Fatalf("expected blocked tool not to execute, got %d calls", tool.callCount) + } +} + func TestPermissionDecisionError(t *testing.T) { t.Parallel() @@ -358,17 +395,19 @@ func TestBuildPermissionAction(t *testing.T) { wantType security.ActionType wantResource string wantTarget string + wantSandbox string wantErr string }{ { name: "bash maps to bash action", input: ToolCallInput{ Name: "bash", - Arguments: []byte(`{"command":"echo hi"}`), + Arguments: []byte(`{"command":"echo hi","workdir":"scripts"}`), }, wantType: security.ActionTypeBash, wantResource: "bash", wantTarget: "echo hi", + wantSandbox: "scripts", }, { name: "read file maps to read action", @@ -379,6 +418,7 @@ func TestBuildPermissionAction(t *testing.T) { wantType: security.ActionTypeRead, wantResource: "filesystem_read_file", wantTarget: "main.go", + wantSandbox: "main.go", }, { name: "grep maps to read action", @@ -389,6 +429,7 @@ func TestBuildPermissionAction(t *testing.T) { wantType: security.ActionTypeRead, wantResource: "filesystem_grep", wantTarget: "internal", + wantSandbox: "internal", }, { name: "glob maps to read action", @@ -399,6 +440,7 @@ func TestBuildPermissionAction(t *testing.T) { wantType: security.ActionTypeRead, wantResource: "filesystem_glob", wantTarget: "cmd", + wantSandbox: "cmd", }, { name: "write file maps to write action", @@ -409,6 +451,7 @@ func TestBuildPermissionAction(t *testing.T) { wantType: security.ActionTypeWrite, wantResource: "filesystem_write_file", wantTarget: "main.go", + wantSandbox: "main.go", }, { name: "webfetch maps to read action", @@ -429,6 +472,7 @@ func TestBuildPermissionAction(t *testing.T) { wantType: security.ActionTypeWrite, wantResource: "filesystem_edit", wantTarget: "main.go", + wantSandbox: "main.go", }, { name: "mcp tool maps to mcp action", @@ -478,6 +522,9 @@ func TestBuildPermissionAction(t *testing.T) { if action.Payload.Target != tt.wantTarget { t.Fatalf("expected target %q, got %q", tt.wantTarget, action.Payload.Target) } + if action.Payload.SandboxTarget != tt.wantSandbox { + t.Fatalf("expected sandbox target %q, got %q", tt.wantSandbox, action.Payload.SandboxTarget) + } }) } } diff --git a/internal/tools/permission_mapper.go b/internal/tools/permission_mapper.go index 04bdf09a..946d1084 100644 --- a/internal/tools/permission_mapper.go +++ b/internal/tools/permission_mapper.go @@ -30,21 +30,29 @@ func buildPermissionAction(input ToolCallInput) (security.Action, error) { action.Payload.Operation = "command" action.Payload.TargetType = security.TargetTypeCommand action.Payload.Target = extractStringArgument(input.Arguments, "command") + action.Payload.SandboxTargetType = security.TargetTypeDirectory + action.Payload.SandboxTarget = extractStringArgument(input.Arguments, "workdir") case "filesystem_read_file": action.Type = security.ActionTypeRead action.Payload.Operation = "read_file" action.Payload.TargetType = security.TargetTypePath action.Payload.Target = extractStringArgument(input.Arguments, "path") + action.Payload.SandboxTargetType = security.TargetTypePath + action.Payload.SandboxTarget = action.Payload.Target case "filesystem_grep": action.Type = security.ActionTypeRead action.Payload.Operation = "grep" action.Payload.TargetType = security.TargetTypeDirectory action.Payload.Target = extractStringArgument(input.Arguments, "dir") + action.Payload.SandboxTargetType = security.TargetTypeDirectory + action.Payload.SandboxTarget = action.Payload.Target case "filesystem_glob": action.Type = security.ActionTypeRead action.Payload.Operation = "glob" action.Payload.TargetType = security.TargetTypeDirectory action.Payload.Target = extractStringArgument(input.Arguments, "dir") + action.Payload.SandboxTargetType = security.TargetTypeDirectory + action.Payload.SandboxTarget = action.Payload.Target case "webfetch": action.Type = security.ActionTypeRead action.Payload.Operation = "fetch" @@ -55,11 +63,15 @@ func buildPermissionAction(input ToolCallInput) (security.Action, error) { action.Payload.Operation = "write_file" action.Payload.TargetType = security.TargetTypePath action.Payload.Target = extractStringArgument(input.Arguments, "path") + action.Payload.SandboxTargetType = security.TargetTypePath + action.Payload.SandboxTarget = action.Payload.Target case "filesystem_edit": action.Type = security.ActionTypeWrite action.Payload.Operation = "edit" action.Payload.TargetType = security.TargetTypePath action.Payload.Target = extractStringArgument(input.Arguments, "path") + action.Payload.SandboxTargetType = security.TargetTypePath + action.Payload.SandboxTarget = action.Payload.Target default: if strings.HasPrefix(strings.ToLower(toolName), "mcp.") { action.Type = security.ActionTypeMCP From 9061c0cea68df2f48b062a2248d8f08d0b06cce7 Mon Sep 17 00:00:00 2001 From: Cai_Tang <106404101+Cai-Tang-www@users.noreply.github.com> Date: Thu, 2 Apr 2026 09:20:21 +0800 Subject: [PATCH 08/23] =?UTF-8?q?app=EF=BC=9A=E5=9C=A8=20bootstrap=20?= =?UTF-8?q?=E4=B8=AD=E5=90=AF=E7=94=A8=20workspace=20sandbox?= MIME-Version: 1.0 Content-Type: text/plain; charset=UTF-8 Content-Transfer-Encoding: 8bit --- internal/app/bootstrap.go | 2 +- internal/app/bootstrap_test.go | 11 +++++++++++ 2 files changed, 12 insertions(+), 1 deletion(-) diff --git a/internal/app/bootstrap.go b/internal/app/bootstrap.go index b6e324f7..1cc077e3 100644 --- a/internal/app/bootstrap.go +++ b/internal/app/bootstrap.go @@ -80,5 +80,5 @@ func buildToolManager(registry *tools.Registry) (tools.Manager, error) { if err != nil { return nil, err } - return tools.NewManager(registry, engine, nil) + return tools.NewManager(registry, engine, security.NewWorkspaceSandbox()) } diff --git a/internal/app/bootstrap_test.go b/internal/app/bootstrap_test.go index 14281848..8b2194bb 100644 --- a/internal/app/bootstrap_test.go +++ b/internal/app/bootstrap_test.go @@ -76,6 +76,7 @@ func TestBuildToolManagerWrapsRegistry(t *testing.T) { registry := tools.NewRegistry() registry.Register(stubToolForBootstrap{name: "bash", content: "ok"}) + workdir := t.TempDir() manager, err := buildToolManager(registry) if err != nil { t.Fatalf("buildToolManager() error = %v", err) @@ -95,6 +96,7 @@ func TestBuildToolManagerWrapsRegistry(t *testing.T) { result, execErr := manager.Execute(context.Background(), tools.ToolCallInput{ Name: "bash", Arguments: []byte(`{"command":"echo hi"}`), + Workdir: workdir, }) if execErr != nil { t.Fatalf("Execute() error = %v", execErr) @@ -102,6 +104,15 @@ func TestBuildToolManagerWrapsRegistry(t *testing.T) { if result.Content != "ok" { t.Fatalf("expected ok result, got %+v", result) } + + _, execErr = manager.Execute(context.Background(), tools.ToolCallInput{ + Name: "bash", + Arguments: []byte(`{"command":"echo hi","workdir":"../outside"}`), + Workdir: workdir, + }) + if execErr == nil { + t.Fatalf("expected sandbox rejection for outside workdir") + } } type stubToolForBootstrap struct { From fc49b29bb01fcc92026ac95609e9a794e7220f22 Mon Sep 17 00:00:00 2001 From: Cai_Tang <106404101+Cai-Tang-www@users.noreply.github.com> Date: Thu, 2 Apr 2026 09:58:44 +0800 Subject: [PATCH 09/23] =?UTF-8?q?feat:=20=E8=A1=A5=E9=BD=90=20workspace=20?= =?UTF-8?q?sandbox=20=E5=88=86=E6=94=AF=E6=B5=8B=E8=AF=95=E8=A6=86?= =?UTF-8?q?=E7=9B=96?= MIME-Version: 1.0 Content-Type: text/plain; charset=UTF-8 Content-Transfer-Encoding: 8bit --- internal/security/workspace_test.go | 505 ++++++++++++++++++++++++++++ 1 file changed, 505 insertions(+) diff --git a/internal/security/workspace_test.go b/internal/security/workspace_test.go index 2a811826..ed061115 100644 --- a/internal/security/workspace_test.go +++ b/internal/security/workspace_test.go @@ -2,6 +2,7 @@ package security import ( "context" + "errors" "os" "path/filepath" "strings" @@ -167,6 +168,510 @@ func TestWorkspaceSandboxCheck(t *testing.T) { } } +func TestWorkspaceSandboxCheckShortCircuits(t *testing.T) { + t.Parallel() + + sandbox := NewWorkspaceSandbox() + root := t.TempDir() + action := fileAction(ActionTypeRead, "filesystem_read_file", "read_file", root, "notes.txt") + + canceledCtx, cancel := context.WithCancel(context.Background()) + cancel() + + err := sandbox.Check(canceledCtx, action) + if !errors.Is(err, context.Canceled) { + t.Fatalf("expected context canceled error, got %v", err) + } + + err = sandbox.Check(context.Background(), Action{ + Type: ActionTypeRead, + Payload: ActionPayload{ + Resource: "filesystem_read_file", + Workdir: root, + }, + }) + if err == nil || !strings.Contains(err.Error(), "tool_name is empty") { + t.Fatalf("expected action validation error, got %v", err) + } +} + +func TestBuildWorkspacePlan(t *testing.T) { + t.Parallel() + + root := t.TempDir() + tests := []struct { + name string + action Action + wantOK bool + wantTarget string + wantErr string + }{ + { + name: "mcp action bypasses workspace sandbox", + action: Action{ + Type: ActionTypeMCP, + Payload: ActionPayload{ + ToolName: "mcp.call", + Resource: "mcp", + Workdir: root, + TargetType: TargetTypeMCP, + }, + }, + wantOK: false, + }, + { + name: "missing root is rejected when sandbox is needed", + action: Action{ + Type: ActionTypeRead, + Payload: ActionPayload{ + ToolName: "filesystem_read_file", + Resource: "filesystem_read_file", + TargetType: TargetTypePath, + Target: "notes.txt", + SandboxTargetType: TargetTypePath, + }, + }, + wantErr: "workspace root is empty", + }, + { + name: "empty path target falls back to tool validation", + action: Action{ + Type: ActionTypeRead, + Payload: ActionPayload{ + ToolName: "filesystem_read_file", + Resource: "filesystem_read_file", + Workdir: root, + TargetType: TargetTypePath, + SandboxTargetType: TargetTypePath, + }, + }, + wantOK: false, + }, + { + name: "directory target defaults to current workspace", + action: Action{ + Type: ActionTypeRead, + Payload: ActionPayload{ + ToolName: "filesystem_grep", + Resource: "filesystem_grep", + Workdir: root, + TargetType: TargetTypeDirectory, + SandboxTargetType: TargetTypeDirectory, + }, + }, + wantOK: true, + wantTarget: ".", + }, + } + + for _, tt := range tests { + tt := tt + t.Run(tt.name, func(t *testing.T) { + t.Parallel() + + plan, ok, err := buildWorkspacePlan(tt.action) + if tt.wantErr != "" { + if err == nil || !strings.Contains(err.Error(), tt.wantErr) { + t.Fatalf("expected error containing %q, got %v", tt.wantErr, err) + } + return + } + + if err != nil { + t.Fatalf("unexpected error: %v", err) + } + if ok != tt.wantOK { + t.Fatalf("ok = %v, want %v", ok, tt.wantOK) + } + if tt.wantTarget != "" && plan.target != tt.wantTarget { + t.Fatalf("target = %q, want %q", plan.target, tt.wantTarget) + } + }) + } +} + +func TestNeedsWorkspaceSandbox(t *testing.T) { + t.Parallel() + + tests := []struct { + name string + typ ActionType + expect bool + }{ + {name: "read is sandboxed", typ: ActionTypeRead, expect: true}, + {name: "write is sandboxed", typ: ActionTypeWrite, expect: true}, + {name: "bash is sandboxed", typ: ActionTypeBash, expect: true}, + {name: "mcp is not sandboxed", typ: ActionTypeMCP, expect: false}, + } + + for _, tt := range tests { + tt := tt + t.Run(tt.name, func(t *testing.T) { + t.Parallel() + if got := needsWorkspaceSandbox(Action{Type: tt.typ}); got != tt.expect { + t.Fatalf("needsWorkspaceSandbox(%q) = %v, want %v", tt.typ, got, tt.expect) + } + }) + } +} + +func TestSandboxTarget(t *testing.T) { + t.Parallel() + + tests := []struct { + name string + action Action + wantTarget string + wantOK bool + }{ + { + name: "bash defaults to current directory", + action: Action{ + Type: ActionTypeBash, + Payload: ActionPayload{ + SandboxTarget: "", + }, + }, + wantTarget: ".", + wantOK: true, + }, + { + name: "bash uses explicit sandbox target", + action: Action{ + Type: ActionTypeBash, + Payload: ActionPayload{ + SandboxTarget: "scripts", + }, + }, + wantTarget: "scripts", + wantOK: true, + }, + { + name: "fallback to target when sandbox target is empty", + action: Action{ + Type: ActionTypeRead, + Payload: ActionPayload{ + TargetType: TargetTypePath, + Target: "main.go", + }, + }, + wantTarget: "main.go", + wantOK: true, + }, + { + name: "directory target defaults to dot", + action: Action{ + Type: ActionTypeRead, + Payload: ActionPayload{ + TargetType: TargetTypeDirectory, + }, + }, + wantTarget: ".", + wantOK: true, + }, + { + name: "empty path target is ignored", + action: Action{ + Type: ActionTypeRead, + Payload: ActionPayload{ + TargetType: TargetTypePath, + }, + }, + wantOK: false, + }, + { + name: "unsupported target type is ignored", + action: Action{ + Type: ActionTypeRead, + Payload: ActionPayload{ + TargetType: TargetTypeURL, + Target: "https://example.com", + }, + }, + wantOK: false, + }, + { + name: "sandbox target takes precedence", + action: Action{ + Type: ActionTypeRead, + Payload: ActionPayload{ + TargetType: TargetTypePath, + Target: "main.go", + SandboxTargetType: TargetTypeDirectory, + SandboxTarget: "docs", + }, + }, + wantTarget: "docs", + wantOK: true, + }, + } + + for _, tt := range tests { + tt := tt + t.Run(tt.name, func(t *testing.T) { + t.Parallel() + gotTarget, gotOK := sandboxTarget(tt.action) + if gotOK != tt.wantOK { + t.Fatalf("ok = %v, want %v", gotOK, tt.wantOK) + } + if gotTarget != tt.wantTarget { + t.Fatalf("target = %q, want %q", gotTarget, tt.wantTarget) + } + }) + } +} + +func TestCanonicalWorkspaceRoot(t *testing.T) { + t.Parallel() + + existing := t.TempDir() + got, err := canonicalWorkspaceRoot(existing) + if err != nil { + t.Fatalf("canonicalWorkspaceRoot(existing) error: %v", err) + } + expectedExisting, err := filepath.Abs(existing) + if err != nil { + t.Fatalf("filepath.Abs(existing): %v", err) + } + if got != filepath.Clean(expectedExisting) { + t.Fatalf("canonicalWorkspaceRoot(existing) = %q, want %q", got, filepath.Clean(expectedExisting)) + } + + missing := filepath.Join(t.TempDir(), "missing", "dir") + got, err = canonicalWorkspaceRoot(missing) + if err != nil { + t.Fatalf("canonicalWorkspaceRoot(missing) error: %v", err) + } + expectedMissing, err := filepath.Abs(missing) + if err != nil { + t.Fatalf("filepath.Abs(missing): %v", err) + } + if got != filepath.Clean(expectedMissing) { + t.Fatalf("canonicalWorkspaceRoot(missing) = %q, want %q", got, filepath.Clean(expectedMissing)) + } + + symlinkRoot := t.TempDir() + targetRoot := filepath.Join(symlinkRoot, "target") + linkRoot := filepath.Join(symlinkRoot, "link") + if err := os.MkdirAll(targetRoot, 0o755); err != nil { + t.Fatalf("mkdir target root: %v", err) + } + if err := os.Symlink(targetRoot, linkRoot); err != nil { + t.Skipf("symlink not supported in this environment: %v", err) + } + got, err = canonicalWorkspaceRoot(linkRoot) + if err != nil { + t.Fatalf("canonicalWorkspaceRoot(link) error: %v", err) + } + expectedLink, err := filepath.Abs(targetRoot) + if err != nil { + t.Fatalf("filepath.Abs(targetRoot): %v", err) + } + if got != filepath.Clean(expectedLink) { + t.Fatalf("canonicalWorkspaceRoot(link) = %q, want %q", got, filepath.Clean(expectedLink)) + } +} + +func TestAbsoluteWorkspaceTarget(t *testing.T) { + t.Parallel() + + root := t.TempDir() + absoluteInside := filepath.Join(root, "abs.txt") + + tests := []struct { + name string + target string + want string + }{ + { + name: "empty target resolves to root", + target: "", + want: root, + }, + { + name: "relative target resolves from root", + target: filepath.Join("dir", "file.txt"), + want: filepath.Join(root, "dir", "file.txt"), + }, + { + name: "absolute target keeps absolute form", + target: absoluteInside, + want: absoluteInside, + }, + } + + for _, tt := range tests { + tt := tt + t.Run(tt.name, func(t *testing.T) { + t.Parallel() + got, err := absoluteWorkspaceTarget(root, tt.target) + if err != nil { + t.Fatalf("absoluteWorkspaceTarget() error: %v", err) + } + wantAbs, err := filepath.Abs(tt.want) + if err != nil { + t.Fatalf("filepath.Abs(%q): %v", tt.want, err) + } + if got != filepath.Clean(wantAbs) { + t.Fatalf("absoluteWorkspaceTarget() = %q, want %q", got, filepath.Clean(wantAbs)) + } + }) + } +} + +func TestEnsureNoSymlinkEscape(t *testing.T) { + t.Parallel() + + tests := []struct { + name string + setup func(t *testing.T) (root string, target string, original string) + expectErr string + }{ + { + name: "root target is allowed", + setup: func(t *testing.T) (string, string, string) { + t.Helper() + root := t.TempDir() + return root, root, "." + }, + }, + { + name: "symlink inside workspace is allowed", + setup: func(t *testing.T) (string, string, string) { + t.Helper() + root := t.TempDir() + realDir := filepath.Join(root, "real") + if err := os.MkdirAll(realDir, 0o755); err != nil { + t.Fatalf("mkdir real dir: %v", err) + } + linkDir := filepath.Join(root, "link") + if err := os.Symlink(realDir, linkDir); err != nil { + t.Skipf("symlink not supported in this environment: %v", err) + } + return root, filepath.Join(root, "link", "new.txt"), filepath.Join("link", "new.txt") + }, + }, + { + name: "broken symlink is rejected", + setup: func(t *testing.T) (string, string, string) { + t.Helper() + root := t.TempDir() + linkDir := filepath.Join(root, "broken") + if err := os.Symlink(filepath.Join(root, "missing"), linkDir); err != nil { + t.Skipf("symlink not supported in this environment: %v", err) + } + return root, filepath.Join(root, "broken", "new.txt"), filepath.Join("broken", "new.txt") + }, + expectErr: "resolve symlink", + }, + } + + for _, tt := range tests { + tt := tt + t.Run(tt.name, func(t *testing.T) { + t.Parallel() + root, target, original := tt.setup(t) + err := ensureNoSymlinkEscape(root, target, original) + if tt.expectErr != "" { + if err == nil || !strings.Contains(err.Error(), tt.expectErr) { + t.Fatalf("expected error containing %q, got %v", tt.expectErr, err) + } + return + } + if err != nil { + t.Fatalf("unexpected error: %v", err) + } + }) + } +} + +func TestSplitRelativePath(t *testing.T) { + t.Parallel() + + tests := []struct { + name string + path string + want []string + }{ + { + name: "dot path returns nil", + path: ".", + want: nil, + }, + { + name: "nested path is split by separator", + path: filepath.Join("a", "b", "c"), + want: []string{"a", "b", "c"}, + }, + } + + for _, tt := range tests { + tt := tt + t.Run(tt.name, func(t *testing.T) { + t.Parallel() + got := splitRelativePath(tt.path) + if len(got) != len(tt.want) { + t.Fatalf("splitRelativePath(%q) len = %d, want %d", tt.path, len(got), len(tt.want)) + } + for i := range got { + if got[i] != tt.want[i] { + t.Fatalf("splitRelativePath(%q)[%d] = %q, want %q", tt.path, i, got[i], tt.want[i]) + } + } + }) + } +} + +func TestIsWithinWorkspace(t *testing.T) { + t.Parallel() + + root := t.TempDir() + inside := filepath.Join(root, "sub", "file.txt") + outside := filepath.Join(t.TempDir(), "outside.txt") + + tests := []struct { + name string + root string + target string + want bool + }{ + { + name: "root itself is inside", + root: root, + target: root, + want: true, + }, + { + name: "child path is inside", + root: root, + target: inside, + want: true, + }, + { + name: "outside path is rejected", + root: root, + target: outside, + want: false, + }, + { + name: "invalid root returns false", + root: "", + target: root, + want: false, + }, + } + + for _, tt := range tests { + tt := tt + t.Run(tt.name, func(t *testing.T) { + t.Parallel() + if got := isWithinWorkspace(tt.root, tt.target); got != tt.want { + t.Fatalf("isWithinWorkspace(%q, %q) = %v, want %v", tt.root, tt.target, got, tt.want) + } + }) + } +} + func fileAction(actionType ActionType, toolName string, operation string, workdir string, target string) Action { return Action{ Type: actionType, From 59d01da7a666fc7dd21a2a48185db395803b877e Mon Sep 17 00:00:00 2001 From: Yumiue <229866007@qq.com> Date: Fri, 3 Apr 2026 09:33:41 +0800 Subject: [PATCH 10/23] =?UTF-8?q?feat:=20=E8=A1=A5=E5=BC=BA=20workspace=20?= =?UTF-8?q?sandbox=20=E8=BE=B9=E7=95=8C=E6=A0=A1=E9=AA=8C=E4=B8=8E?= =?UTF-8?q?=E6=96=87=E6=A1=A3=E8=AF=B4=E6=98=8E?= MIME-Version: 1.0 Content-Type: text/plain; charset=UTF-8 Content-Transfer-Encoding: 8bit --- internal/security/types.go | 22 ++++++++----- internal/security/workspace.go | 59 ++++++++++++++++++++++++++++++++-- 2 files changed, 70 insertions(+), 11 deletions(-) diff --git a/internal/security/types.go b/internal/security/types.go index 820d74fe..50da521b 100644 --- a/internal/security/types.go +++ b/internal/security/types.go @@ -58,15 +58,21 @@ const ( // ActionPayload is the normalized structured context used by policy and sandbox. type ActionPayload struct { - ToolName string - Resource string - Operation string - SessionID string - Workdir string - TargetType TargetType - Target string + ToolName string + Resource string + Operation string + SessionID string + Workdir string + TargetType TargetType + Target string + // SandboxTargetType is the target kind used specifically for workspace boundary + // checks. It falls back to TargetType when unset so callers can keep permission + // display metadata separate from the path actually validated by the sandbox. SandboxTargetType TargetType - SandboxTarget string + // SandboxTarget is the concrete path/value used specifically for workspace + // validation. It falls back to Target when unset. For example, bash validates + // its requested workdir here while Target still contains the shell command. + SandboxTarget string } // Action is the unified security input for one tool execution request. diff --git a/internal/security/workspace.go b/internal/security/workspace.go index 65335f35..6b30a824 100644 --- a/internal/security/workspace.go +++ b/internal/security/workspace.go @@ -42,6 +42,9 @@ type workspacePlan struct { target string } +// buildWorkspacePlan extracts the workspace root and the sandbox-specific target +// that should be validated. Actions that do not touch local workspace paths are +// ignored here and continue through the permission pipeline unchanged. func buildWorkspacePlan(action Action) (workspacePlan, bool, error) { if !needsWorkspaceSandbox(action) { return workspacePlan{}, false, nil @@ -72,6 +75,10 @@ func needsWorkspaceSandbox(action Action) bool { } } +// sandboxTarget returns the path-like value that should be checked against the +// workspace boundary. This is intentionally separate from ActionPayload.Target, +// because some tools expose one value to policy while validating another one for +// local filesystem access. func sandboxTarget(action Action) (string, bool) { if action.Type == ActionTypeBash { target := strings.TrimSpace(action.Payload.SandboxTarget) @@ -107,6 +114,9 @@ func sandboxTarget(action Action) (string, bool) { } } +// validateWorkspacePlan resolves the canonical workspace root, expands the +// requested target to an absolute path, verifies it stays inside the workspace, +// and then checks every existing path segment for symlink escape. func validateWorkspacePlan(plan workspacePlan) error { root, err := canonicalWorkspaceRoot(plan.root) if err != nil { @@ -124,23 +134,34 @@ func validateWorkspacePlan(plan workspacePlan) error { return ensureNoSymlinkEscape(root, target, plan.target) } +// canonicalWorkspaceRoot resolves the configured workspace root to a canonical +// directory path. The root must already exist so validation cannot silently rely +// on a string-only path that may later resolve through a symlinked parent. func canonicalWorkspaceRoot(root string) (string, error) { absoluteRoot, err := filepath.Abs(strings.TrimSpace(root)) if err != nil { return "", fmt.Errorf("security: resolve workspace root: %w", err) } + info, err := os.Stat(absoluteRoot) + if err != nil { + return "", fmt.Errorf("security: resolve workspace root: %w", err) + } + if !info.IsDir() { + return "", fmt.Errorf("security: workspace root %q is not a directory", root) + } + canonicalRoot, err := filepath.EvalSymlinks(absoluteRoot) if err != nil { - if errors.Is(err, os.ErrNotExist) { - return filepath.Clean(absoluteRoot), nil - } return "", fmt.Errorf("security: resolve workspace root: %w", err) } return filepath.Clean(canonicalRoot), nil } +// absoluteWorkspaceTarget expands a workspace-relative target into an absolute +// path and rejects Windows volume changes up front so later boundary checks do +// not depend on platform-specific filepath.Rel behavior. func absoluteWorkspaceTarget(root string, target string) (string, error) { trimmedTarget := strings.TrimSpace(target) if trimmedTarget == "" { @@ -155,9 +176,23 @@ func absoluteWorkspaceTarget(root string, target string) (string, error) { return "", fmt.Errorf("security: resolve workspace target %q: %w", target, err) } + if err := validateTargetVolume(root, absoluteTarget); err != nil { + return "", err + } + return filepath.Clean(absoluteTarget), nil } +// ensureNoSymlinkEscape walks each existing segment from root to target and +// ensures that any encountered symlink still resolves inside the workspace. +// +// This blocks common symlink escape attempts such as placing a link inside the +// workspace that points to a file or directory outside the workspace tree. +// +// Known limitation: this validation is still subject to TOCTOU races between the +// sandbox check and the later filesystem operation in the executor. Eliminating +// that window requires executor-level changes such as descriptor-based access or +// no-follow file opening, which are outside this lightweight sandbox layer. func ensureNoSymlinkEscape(root string, target string, original string) error { relativeTarget, err := filepath.Rel(root, target) if err != nil { @@ -202,6 +237,24 @@ func ensureNoSymlinkEscape(root string, target string, original string) error { return nil } +func validateTargetVolume(root string, target string) error { + rootVolume := normalizeVolumeName(root) + targetVolume := normalizeVolumeName(target) + if rootVolume == "" || targetVolume == "" { + return nil + } + if !strings.EqualFold(rootVolume, targetVolume) { + return fmt.Errorf("security: path %q is on different volume than workspace root", target) + } + return nil +} + +func normalizeVolumeName(path string) string { + volume := strings.TrimSpace(filepath.VolumeName(filepath.Clean(path))) + volume = strings.TrimPrefix(volume, `\\?\`) + return strings.ToLower(volume) +} + func splitRelativePath(path string) []string { cleanPath := filepath.Clean(path) if cleanPath == "." { From d07925a3c0df0c90ae621cd3cd6155d2b7350de8 Mon Sep 17 00:00:00 2001 From: Cai_Tang <106404101+Cai-Tang-www@users.noreply.github.com> Date: Thu, 2 Apr 2026 16:18:42 +0800 Subject: [PATCH 11/23] =?UTF-8?q?feat:=20=E8=A1=A5=E5=BC=BA=20workspace=20?= =?UTF-8?q?sandbox=20=E8=BE=B9=E7=95=8C=E6=A0=A1=E9=AA=8C=E4=B8=8E?= =?UTF-8?q?=E6=96=87=E6=A1=A3=E8=AF=B4=E6=98=8E?= MIME-Version: 1.0 Content-Type: text/plain; charset=UTF-8 Content-Transfer-Encoding: 8bit --- internal/security/workspace_test.go | 78 +++++++++++++++++++++++++---- 1 file changed, 69 insertions(+), 9 deletions(-) diff --git a/internal/security/workspace_test.go b/internal/security/workspace_test.go index ed061115..e55519d8 100644 --- a/internal/security/workspace_test.go +++ b/internal/security/workspace_test.go @@ -5,6 +5,7 @@ import ( "errors" "os" "path/filepath" + "runtime" "strings" "testing" ) @@ -438,16 +439,16 @@ func TestCanonicalWorkspaceRoot(t *testing.T) { } missing := filepath.Join(t.TempDir(), "missing", "dir") - got, err = canonicalWorkspaceRoot(missing) - if err != nil { - t.Fatalf("canonicalWorkspaceRoot(missing) error: %v", err) - } - expectedMissing, err := filepath.Abs(missing) - if err != nil { - t.Fatalf("filepath.Abs(missing): %v", err) + _, err = canonicalWorkspaceRoot(missing) + if err == nil || !strings.Contains(err.Error(), "resolve workspace root") { + t.Fatalf("expected missing root error, got %v", err) } - if got != filepath.Clean(expectedMissing) { - t.Fatalf("canonicalWorkspaceRoot(missing) = %q, want %q", got, filepath.Clean(expectedMissing)) + + notDirRoot := filepath.Join(t.TempDir(), "file.txt") + mustWriteWorkspaceFile(t, notDirRoot, "content") + _, err = canonicalWorkspaceRoot(notDirRoot) + if err == nil || !strings.Contains(err.Error(), "is not a directory") { + t.Fatalf("expected not-a-directory error, got %v", err) } symlinkRoot := t.TempDir() @@ -519,6 +520,65 @@ func TestAbsoluteWorkspaceTarget(t *testing.T) { } } +func TestValidateTargetVolume(t *testing.T) { + t.Parallel() + + if runtime.GOOS != "windows" { + t.Skip("volume validation is Windows-specific") + } + + root := `C:\workspace` + inside := `C:\workspace\file.txt` + outside := `D:\secret.txt` + + if err := validateTargetVolume(root, inside); err != nil { + t.Fatalf("validateTargetVolume(inside) error: %v", err) + } + if err := validateTargetVolume(root, outside); err == nil || !strings.Contains(err.Error(), "different volume") { + t.Fatalf("expected different volume error, got %v", err) + } +} + +func TestNormalizeVolumeName(t *testing.T) { + t.Parallel() + + if runtime.GOOS != "windows" { + t.Skip("volume normalization is Windows-specific") + } + + tests := []struct { + name string + path string + want string + }{ + { + name: "drive letter is normalized", + path: `C:\workspace`, + want: `c:`, + }, + { + name: "extended path prefix is removed", + path: `\\?\C:\workspace`, + want: `c:`, + }, + { + name: "unc share is normalized", + path: `\\server\share\dir`, + want: `\\server\share`, + }, + } + + for _, tt := range tests { + tt := tt + t.Run(tt.name, func(t *testing.T) { + t.Parallel() + if got := normalizeVolumeName(tt.path); got != tt.want { + t.Fatalf("normalizeVolumeName(%q) = %q, want %q", tt.path, got, tt.want) + } + }) + } +} + func TestEnsureNoSymlinkEscape(t *testing.T) { t.Parallel() From 63c3dd671c7380783ecae843abe8b6d6919b17b9 Mon Sep 17 00:00:00 2001 From: Cai_Tang <106404101+Cai-Tang-www@users.noreply.github.com> Date: Thu, 2 Apr 2026 16:41:26 +0800 Subject: [PATCH 12/23] =?UTF-8?q?feat:=20=E4=BC=98=E5=8C=96=20workspace=20?= =?UTF-8?q?sandbox=20=E6=A0=B9=E8=B7=AF=E5=BE=84=E7=BC=93=E5=AD=98?= =?UTF-8?q?=E4=B8=8E=E8=A7=A3=E6=9E=90=E6=80=A7=E8=83=BD?= MIME-Version: 1.0 Content-Type: text/plain; charset=UTF-8 Content-Transfer-Encoding: 8bit --- internal/security/workspace.go | 145 ++++++++++++++++++---------- internal/security/workspace_test.go | 94 ++++++++++++++++-- 2 files changed, 183 insertions(+), 56 deletions(-) diff --git a/internal/security/workspace.go b/internal/security/workspace.go index 6b30a824..a0c6198a 100644 --- a/internal/security/workspace.go +++ b/internal/security/workspace.go @@ -6,11 +6,15 @@ import ( "fmt" "os" "path/filepath" + "runtime" "strings" + "sync" ) // WorkspaceSandbox enforces workspace-relative path boundaries for tool actions. -type WorkspaceSandbox struct{} +type WorkspaceSandbox struct { + canonicalRoots sync.Map +} // NewWorkspaceSandbox creates a sandbox that blocks traversal and symlink escape. func NewWorkspaceSandbox() *WorkspaceSandbox { @@ -34,7 +38,7 @@ func (s *WorkspaceSandbox) Check(ctx context.Context, action Action) error { return nil } - return validateWorkspacePlan(plan) + return s.validateWorkspacePlan(plan) } type workspacePlan struct { @@ -116,9 +120,9 @@ func sandboxTarget(action Action) (string, bool) { // validateWorkspacePlan resolves the canonical workspace root, expands the // requested target to an absolute path, verifies it stays inside the workspace, -// and then checks every existing path segment for symlink escape. -func validateWorkspacePlan(plan workspacePlan) error { - root, err := canonicalWorkspaceRoot(plan.root) +// and then checks the nearest existing path ancestor for symlink escape. +func (s *WorkspaceSandbox) validateWorkspacePlan(plan workspacePlan) error { + root, err := s.canonicalWorkspaceRoot(plan.root) if err != nil { return err } @@ -135,28 +139,47 @@ func validateWorkspacePlan(plan workspacePlan) error { } // canonicalWorkspaceRoot resolves the configured workspace root to a canonical -// directory path. The root must already exist so validation cannot silently rely -// on a string-only path that may later resolve through a symlinked parent. -func canonicalWorkspaceRoot(root string) (string, error) { +// directory path and caches non-symlink workspace roots for repeated tool calls. +func (s *WorkspaceSandbox) canonicalWorkspaceRoot(root string) (string, error) { absoluteRoot, err := filepath.Abs(strings.TrimSpace(root)) if err != nil { return "", fmt.Errorf("security: resolve workspace root: %w", err) } + cacheKey := cleanedPathKey(absoluteRoot) + if cached, ok := s.canonicalRoots.Load(cacheKey); ok { + return cached.(string), nil + } + + canonicalRoot, cacheable, err := resolveCanonicalWorkspaceRoot(cacheKey) + if err != nil { + return "", err + } + if cacheable { + s.canonicalRoots.Store(cacheKey, canonicalRoot) + } + return canonicalRoot, nil +} +// resolveCanonicalWorkspaceRoot resolves the configured workspace root to a +// canonical directory path. The root must already exist so validation cannot +// silently rely on a string-only path that may later resolve through a +// symlinked parent. +func resolveCanonicalWorkspaceRoot(absoluteRoot string) (string, bool, error) { info, err := os.Stat(absoluteRoot) if err != nil { - return "", fmt.Errorf("security: resolve workspace root: %w", err) + return "", false, fmt.Errorf("security: resolve workspace root: %w", err) } if !info.IsDir() { - return "", fmt.Errorf("security: workspace root %q is not a directory", root) + return "", false, fmt.Errorf("security: workspace root %q is not a directory", absoluteRoot) } canonicalRoot, err := filepath.EvalSymlinks(absoluteRoot) if err != nil { - return "", fmt.Errorf("security: resolve workspace root: %w", err) + return "", false, fmt.Errorf("security: resolve workspace root: %w", err) } - return filepath.Clean(canonicalRoot), nil + cleanedCanonical := cleanedPathKey(canonicalRoot) + return cleanedCanonical, samePathKey(absoluteRoot, cleanedCanonical), nil } // absoluteWorkspaceTarget expands a workspace-relative target into an absolute @@ -183,57 +206,38 @@ func absoluteWorkspaceTarget(root string, target string) (string, error) { return filepath.Clean(absoluteTarget), nil } -// ensureNoSymlinkEscape walks each existing segment from root to target and -// ensures that any encountered symlink still resolves inside the workspace. +// ensureNoSymlinkEscape resolves the nearest existing path on the way to target +// and ensures that any symlink in the existing prefix still resolves inside the +// workspace. // // This blocks common symlink escape attempts such as placing a link inside the -// workspace that points to a file or directory outside the workspace tree. +// workspace that points to a file or directory outside the workspace tree while +// avoiding a path-by-path stat on every component. // // Known limitation: this validation is still subject to TOCTOU races between the // sandbox check and the later filesystem operation in the executor. Eliminating // that window requires executor-level changes such as descriptor-based access or // no-follow file opening, which are outside this lightweight sandbox layer. func ensureNoSymlinkEscape(root string, target string, original string) error { - relativeTarget, err := filepath.Rel(root, target) + existingPath, err := nearestExistingPath(root, target) if err != nil { - return fmt.Errorf("security: compare workspace target %q: %w", original, err) + return err } - - cleanRelative := filepath.Clean(relativeTarget) - if cleanRelative == "." { + if existingPath == root { return nil } - current := root - for _, segment := range splitRelativePath(cleanRelative) { - next := filepath.Join(current, segment) - info, err := os.Lstat(next) - if err != nil { - if errors.Is(err, os.ErrNotExist) { - return nil - } - return fmt.Errorf("security: inspect path %q: %w", next, err) - } - - if info.Mode()&os.ModeSymlink != 0 { - resolved, err := filepath.EvalSymlinks(next) - if err != nil { - return fmt.Errorf("security: resolve symlink %q: %w", next, err) - } - resolved, err = filepath.Abs(resolved) - if err != nil { - return fmt.Errorf("security: resolve symlink %q: %w", next, err) - } - if !isWithinWorkspace(root, resolved) { - return fmt.Errorf("security: path %q escapes workspace root via symlink", original) - } - current = filepath.Clean(resolved) - continue - } - - current = next + resolved, err := filepath.EvalSymlinks(existingPath) + if err != nil { + return fmt.Errorf("security: resolve symlink %q: %w", existingPath, err) + } + resolved, err = filepath.Abs(resolved) + if err != nil { + return fmt.Errorf("security: resolve symlink %q: %w", existingPath, err) + } + if !isWithinWorkspace(root, resolved) { + return fmt.Errorf("security: path %q escapes workspace root via symlink", original) } - return nil } @@ -250,11 +254,52 @@ func validateTargetVolume(root string, target string) error { } func normalizeVolumeName(path string) string { - volume := strings.TrimSpace(filepath.VolumeName(filepath.Clean(path))) + volume := strings.TrimSpace(filepath.VolumeName(cleanedPathKey(path))) volume = strings.TrimPrefix(volume, `\\?\`) return strings.ToLower(volume) } +func nearestExistingPath(root string, target string) (string, error) { + current := cleanedPathKey(target) + root = cleanedPathKey(root) + for { + info, err := os.Lstat(current) + if err == nil { + if info.Mode()&os.ModeSymlink != 0 || current != root { + return current, nil + } + return root, nil + } + if !errors.Is(err, os.ErrNotExist) { + return "", fmt.Errorf("security: inspect path %q: %w", current, err) + } + if samePathKey(current, root) { + return root, nil + } + + parent := filepath.Dir(current) + if samePathKey(parent, current) { + return "", fmt.Errorf("security: compare workspace target %q: reached filesystem root", target) + } + current = parent + } +} + +func cleanedPathKey(path string) string { + cleaned := filepath.Clean(path) + if runtime.GOOS == "windows" { + return strings.ToLower(cleaned) + } + return cleaned +} + +func samePathKey(a string, b string) bool { + if runtime.GOOS == "windows" { + return strings.EqualFold(cleanedPathKey(a), cleanedPathKey(b)) + } + return cleanedPathKey(a) == cleanedPathKey(b) +} + func splitRelativePath(path string) []string { cleanPath := filepath.Clean(path) if cleanPath == "." { diff --git a/internal/security/workspace_test.go b/internal/security/workspace_test.go index e55519d8..4cb75c14 100644 --- a/internal/security/workspace_test.go +++ b/internal/security/workspace_test.go @@ -425,8 +425,9 @@ func TestSandboxTarget(t *testing.T) { func TestCanonicalWorkspaceRoot(t *testing.T) { t.Parallel() + sandbox := NewWorkspaceSandbox() existing := t.TempDir() - got, err := canonicalWorkspaceRoot(existing) + got, err := sandbox.canonicalWorkspaceRoot(existing) if err != nil { t.Fatalf("canonicalWorkspaceRoot(existing) error: %v", err) } @@ -434,19 +435,22 @@ func TestCanonicalWorkspaceRoot(t *testing.T) { if err != nil { t.Fatalf("filepath.Abs(existing): %v", err) } - if got != filepath.Clean(expectedExisting) { + if !samePathKey(got, expectedExisting) { t.Fatalf("canonicalWorkspaceRoot(existing) = %q, want %q", got, filepath.Clean(expectedExisting)) } + if _, ok := sandbox.canonicalRoots.Load(cleanedPathKey(existing)); !ok { + t.Fatalf("expected canonical root cache entry for %q", existing) + } missing := filepath.Join(t.TempDir(), "missing", "dir") - _, err = canonicalWorkspaceRoot(missing) + _, err = sandbox.canonicalWorkspaceRoot(missing) if err == nil || !strings.Contains(err.Error(), "resolve workspace root") { t.Fatalf("expected missing root error, got %v", err) } notDirRoot := filepath.Join(t.TempDir(), "file.txt") mustWriteWorkspaceFile(t, notDirRoot, "content") - _, err = canonicalWorkspaceRoot(notDirRoot) + _, err = sandbox.canonicalWorkspaceRoot(notDirRoot) if err == nil || !strings.Contains(err.Error(), "is not a directory") { t.Fatalf("expected not-a-directory error, got %v", err) } @@ -460,7 +464,7 @@ func TestCanonicalWorkspaceRoot(t *testing.T) { if err := os.Symlink(targetRoot, linkRoot); err != nil { t.Skipf("symlink not supported in this environment: %v", err) } - got, err = canonicalWorkspaceRoot(linkRoot) + got, err = sandbox.canonicalWorkspaceRoot(linkRoot) if err != nil { t.Fatalf("canonicalWorkspaceRoot(link) error: %v", err) } @@ -468,9 +472,12 @@ func TestCanonicalWorkspaceRoot(t *testing.T) { if err != nil { t.Fatalf("filepath.Abs(targetRoot): %v", err) } - if got != filepath.Clean(expectedLink) { + if !samePathKey(got, expectedLink) { t.Fatalf("canonicalWorkspaceRoot(link) = %q, want %q", got, filepath.Clean(expectedLink)) } + if _, ok := sandbox.canonicalRoots.Load(cleanedPathKey(linkRoot)); ok { + t.Fatalf("did not expect symlinked root %q to be cached", linkRoot) + } } func TestAbsoluteWorkspaceTarget(t *testing.T) { @@ -645,6 +652,81 @@ func TestEnsureNoSymlinkEscape(t *testing.T) { } } +func TestNearestExistingPath(t *testing.T) { + t.Parallel() + + tests := []struct { + name string + setup func(t *testing.T) (root string, target string) + expect func(root string, target string) string + expectErr string + }{ + { + name: "returns root for missing nested path", + setup: func(t *testing.T) (string, string) { + t.Helper() + root := t.TempDir() + return root, filepath.Join(root, "missing", "file.txt") + }, + expect: func(root string, target string) string { + return cleanedPathKey(root) + }, + }, + { + name: "returns nearest existing ancestor", + setup: func(t *testing.T) (string, string) { + t.Helper() + root := t.TempDir() + dir := filepath.Join(root, "dir") + if err := os.MkdirAll(dir, 0o755); err != nil { + t.Fatalf("mkdir dir: %v", err) + } + return root, filepath.Join(dir, "missing", "file.txt") + }, + expect: func(root string, target string) string { + return cleanedPathKey(filepath.Join(root, "dir")) + }, + }, + { + name: "returns broken symlink ancestor", + setup: func(t *testing.T) (string, string) { + t.Helper() + root := t.TempDir() + link := filepath.Join(root, "broken") + if err := os.Symlink(filepath.Join(root, "missing"), link); err != nil { + t.Skipf("symlink not supported in this environment: %v", err) + } + return root, filepath.Join(link, "child.txt") + }, + expect: func(root string, target string) string { + return cleanedPathKey(filepath.Join(root, "broken")) + }, + }, + } + + for _, tt := range tests { + tt := tt + t.Run(tt.name, func(t *testing.T) { + t.Parallel() + root, target := tt.setup(t) + got, err := nearestExistingPath(root, target) + if tt.expectErr != "" { + if err == nil || !strings.Contains(err.Error(), tt.expectErr) { + t.Fatalf("expected error containing %q, got %v", tt.expectErr, err) + } + return + } + if err != nil { + t.Fatalf("nearestExistingPath() error: %v", err) + } + want := tt.expect(root, target) + if got != want { + t.Fatalf("nearestExistingPath() = %q, want %q", got, want) + } + }) + } +} + func TestSplitRelativePath(t *testing.T) { t.Parallel() From 51749192aee05ade7b5e08382d014c0d57f7ea9a Mon Sep 17 00:00:00 2001 From: creatang Date: Thu, 2 Apr 2026 09:31:53 +0800 Subject: [PATCH 13/23] fix(tui): separate activity from transcript --- internal/tui/app.go | 1 + internal/tui/commands.go | 6 ++ internal/tui/update.go | 166 +++++++++++++++++++++++++++--------- internal/tui/update_test.go | 165 ++++++++++++++++++++++++++++++++--- internal/tui/view.go | 50 +++++++++++ 5 files changed, 338 insertions(+), 50 deletions(-) diff --git a/internal/tui/app.go b/internal/tui/app.go index 830a6678..08cd4bff 100644 --- a/internal/tui/app.go +++ b/internal/tui/app.go @@ -30,6 +30,7 @@ type App struct { transcript viewport.Model input textarea.Model activeMessages []provider.Message + activities []activityEntry fileCandidates []string focus panel width int diff --git a/internal/tui/commands.go b/internal/tui/commands.go index 094d4dc0..0cf8989b 100644 --- a/internal/tui/commands.go +++ b/internal/tui/commands.go @@ -39,6 +39,8 @@ const ( sidebarTitle = "Sessions" sidebarFilterHint = "Type / to search" sidebarOpenHint = "Enter to open" + activityTitle = "Activity" + activitySubtitle = "Latest execution events" draftSessionTitle = "Draft" emptyConversationText = "No conversation yet.\nAsk NeoCode to inspect or change code, or type /help to browse local commands." @@ -64,8 +66,12 @@ const ( focusLabelSessions = "Sessions" focusLabelTranscript = "Transcript" + focusLabelActivity = "Activity" focusLabelComposer = "Composer" + activityPreviewEntries = 3 + maxActivityEntries = 64 + messageTagUser = "[ YOU ]" messageTagAgent = "[ NEO ]" messageTagTool = "[ TOOL ]" diff --git a/internal/tui/update.go b/internal/tui/update.go index a37e44d6..6b5b5fc1 100644 --- a/internal/tui/update.go +++ b/internal/tui/update.go @@ -5,6 +5,7 @@ import ( "errors" "fmt" "strings" + "time" "github.com/charmbracelet/bubbles/key" "github.com/charmbracelet/bubbles/list" @@ -33,6 +34,7 @@ const ( composerMinHeight = 1 composerMaxHeight = 5 composerPromptWidth = 2 + mouseWheelStepLines = 3 ) func (a App) Update(msg tea.Msg) (tea.Model, tea.Cmd) { @@ -50,10 +52,12 @@ func (a App) Update(msg tea.Msg) (tea.Model, tea.Cmd) { a.resizeComponents() return a, tea.Batch(cmds...) case RuntimeMsg: - a.handleRuntimeEvent(typed.Event) + transcriptDirty := a.handleRuntimeEvent(typed.Event) _ = a.refreshSessions() a.syncActiveSessionTitle() - a.rebuildTranscript() + if transcriptDirty { + a.rebuildTranscript() + } cmds = append(cmds, ListenForRuntimeEvent(a.runtime.Events())) return a, tea.Batch(cmds...) case RuntimeClosedMsg: @@ -61,7 +65,6 @@ func (a App) Update(msg tea.Msg) (tea.Model, tea.Cmd) { if strings.TrimSpace(a.state.StatusText) == "" { a.state.StatusText = statusRuntimeClosed } - a.rebuildTranscript() return a, tea.Batch(cmds...) case runFinishedMsg: if typed.err != nil { @@ -74,12 +77,10 @@ func (a App) Update(msg tea.Msg) (tea.Model, tea.Cmd) { } else { a.state.ExecutionError = typed.err.Error() a.state.StatusText = typed.err.Error() - a.appendInlineMessage(roleError, typed.err.Error()) } } _ = a.refreshSessions() a.syncActiveSessionTitle() - a.rebuildTranscript() return a, tea.Batch(cmds...) case compactFinishedMsg: // compact feedback should primarily come from runtime compact_done/error events. @@ -105,7 +106,7 @@ func (a App) Update(msg tea.Msg) (tea.Model, tea.Cmd) { if typed.err != nil { a.state.ExecutionError = typed.err.Error() a.state.StatusText = typed.err.Error() - a.appendInlineMessage(roleError, typed.err.Error()) + a.appendActivity("command", "Local command failed", typed.err.Error(), true) } else { a.state.ExecutionError = "" a.state.StatusText = typed.notice @@ -114,42 +115,41 @@ func (a App) Update(msg tea.Msg) (tea.Model, tea.Cmd) { if err := a.refreshProviderPicker(); err != nil { a.state.ExecutionError = err.Error() a.state.StatusText = err.Error() - a.appendInlineMessage(roleError, err.Error()) + a.appendActivity("system", "Failed to refresh providers", err.Error(), true) } if err := a.refreshModelPicker(); err != nil { a.state.ExecutionError = err.Error() a.state.StatusText = err.Error() - a.appendInlineMessage(roleError, err.Error()) + a.appendActivity("system", "Failed to refresh models", err.Error(), true) } else { a.selectCurrentProvider(cfg.SelectedProvider) a.selectCurrentModel(cfg.CurrentModel) } - a.appendInlineMessage(roleSystem, typed.notice) + a.appendActivity("command", typed.notice, "", false) } - a.rebuildTranscript() - a.transcript.GotoBottom() return a, tea.Batch(cmds...) case workspaceCommandResultMsg: if typed.command == "" && typed.err != nil { a.state.ExecutionError = typed.err.Error() a.state.StatusText = typed.err.Error() - a.appendInlineMessage(roleError, typed.err.Error()) - a.rebuildTranscript() - a.transcript.GotoBottom() + a.appendActivity("command", "Workspace command failed", typed.err.Error(), true) return a, tea.Batch(cmds...) } result := formatWorkspaceCommandResult(typed.command, typed.output, typed.err) - a.activeMessages = append(a.activeMessages, provider.Message{Role: roleTool, Content: result}) if typed.err != nil { a.state.ExecutionError = typed.err.Error() a.state.StatusText = fmt.Sprintf("Command failed: %s", typed.command) + a.appendActivity("command", "Command failed", result, true) } else { a.state.ExecutionError = "" a.state.StatusText = statusCommandDone + a.appendActivity("command", "Command finished", result, false) } - a.rebuildTranscript() - a.transcript.GotoBottom() return a, tea.Batch(cmds...) + case tea.MouseMsg: + if a.handleTranscriptMouse(typed) { + return a, tea.Batch(cmds...) + } case tea.KeyMsg: if key.Matches(typed, a.keys.Quit) { return a, tea.Quit @@ -164,7 +164,6 @@ func (a App) Update(msg tea.Msg) (tea.Model, tea.Cmd) { if a.runtime.CancelActiveRun() { a.state.StatusText = statusCanceling } - a.rebuildTranscript() return a, tea.Batch(cmds...) } if a.state.ActivePicker != pickerNone { @@ -197,11 +196,10 @@ func (a App) Update(msg tea.Msg) (tea.Model, tea.Cmd) { if err := a.activateSelectedSession(); err != nil { a.state.StatusText = err.Error() a.state.ExecutionError = err.Error() - a.appendInlineMessage(roleError, err.Error()) + a.appendActivity("system", "Failed to open session", err.Error(), true) } a.focus = panelInput a.applyFocus() - a.rebuildTranscript() return a, tea.Batch(cmds...) } var cmd tea.Cmd @@ -243,8 +241,7 @@ func (a App) updateInputPanel(msg tea.Msg, typed tea.KeyMsg, cmds []tea.Cmd) (te if err := a.refreshProviderPicker(); err != nil { a.state.ExecutionError = err.Error() a.state.StatusText = err.Error() - a.appendInlineMessage(roleError, err.Error()) - a.rebuildTranscript() + a.appendActivity("system", "Failed to refresh providers", err.Error(), true) return a, tea.Batch(cmds...) } a.openProviderPicker() @@ -253,8 +250,7 @@ func (a App) updateInputPanel(msg tea.Msg, typed tea.KeyMsg, cmds []tea.Cmd) (te if err := a.refreshModelPicker(); err != nil { a.state.ExecutionError = err.Error() a.state.StatusText = err.Error() - a.appendInlineMessage(roleError, err.Error()) - a.rebuildTranscript() + a.appendActivity("system", "Failed to refresh models", err.Error(), true) return a, tea.Batch(cmds...) } a.openModelPicker() @@ -271,20 +267,18 @@ func (a App) updateInputPanel(msg tea.Msg, typed tea.KeyMsg, cmds []tea.Cmd) (te if err != nil { a.state.ExecutionError = err.Error() a.state.StatusText = err.Error() - a.appendInlineMessage(roleError, err.Error()) - a.rebuildTranscript() - a.transcript.GotoBottom() + a.appendActivity("command", "Invalid workspace command", err.Error(), true) return a, tea.Batch(cmds...) } + a.activities = nil a.state.StatusText = statusRunningCommand a.state.ExecutionError = "" - a.appendInlineMessage(roleEvent, "Running command: "+command) - a.rebuildTranscript() - a.transcript.GotoBottom() + a.appendActivity("command", "Running command", command, false) cmds = append(cmds, runWorkspaceCommand(a.configManager, input)) return a, tea.Batch(cmds...) } + a.activities = nil a.state.IsAgentRunning = true a.state.StreamingReply = false a.state.ExecutionError = "" @@ -377,6 +371,7 @@ func (a *App) refreshSessions() error { func (a *App) refreshMessages() error { if strings.TrimSpace(a.state.ActiveSessionID) == "" { a.activeMessages = nil + a.activities = nil return nil } @@ -386,6 +381,7 @@ func (a *App) refreshMessages() error { } a.activeMessages = session.Messages + a.activities = nil a.state.ActiveSessionTitle = session.Title return nil } @@ -430,45 +426,59 @@ func (a *App) syncConfigState(cfg config.Config) { a.state.CurrentWorkdir = cfg.Workdir } -func (a *App) handleRuntimeEvent(event agentruntime.RuntimeEvent) { +func (a *App) handleRuntimeEvent(event agentruntime.RuntimeEvent) bool { if a.state.ActiveSessionID == "" { a.state.ActiveSessionID = event.SessionID } + transcriptDirty := false + switch event.Type { case agentruntime.EventUserMessage: a.state.StatusText = statusThinking a.state.StreamingReply = false a.state.CurrentTool = "" a.state.ExecutionError = "" + case agentruntime.EventToolCallThinking: + if payload, ok := event.Payload.(string); ok && strings.TrimSpace(payload) != "" { + a.state.CurrentTool = payload + a.appendActivity("tool", "Planning tool call", payload, false) + } case agentruntime.EventToolStart: a.state.StatusText = statusRunningTool a.state.StreamingReply = false if payload, ok := event.Payload.(provider.ToolCall); ok { a.state.CurrentTool = payload.Name - a.appendInlineMessage(roleEvent, "Running tool: "+payload.Name+"...") + a.appendActivity("tool", "Running tool", payload.Name, false) } case agentruntime.EventToolResult: a.state.StreamingReply = false a.state.CurrentTool = "" if payload, ok := event.Payload.(tools.ToolResult); ok { + a.activeMessages = append(a.activeMessages, provider.Message{ + Role: roleTool, + Content: payload.Content, + IsError: payload.IsError, + }) + transcriptDirty = true if payload.IsError { a.state.ExecutionError = payload.Content a.state.StatusText = statusToolError - a.appendInlineMessage(roleError, preview(payload.Content, 88, 4)) + a.appendActivity("tool", "Tool error", preview(payload.Content, 88, 4), true) } else if strings.TrimSpace(a.state.ExecutionError) == "" { a.state.StatusText = statusToolFinished - a.appendInlineMessage(roleEvent, "Completed tool: "+payload.Name) + a.appendActivity("tool", "Completed tool", payload.Name, false) } } case agentruntime.EventAgentChunk: if payload, ok := event.Payload.(string); ok { a.appendAssistantChunk(payload) + transcriptDirty = true } case agentruntime.EventToolChunk: if payload, ok := event.Payload.(string); ok && strings.TrimSpace(payload) != "" { a.state.StatusText = statusRunningTool - a.appendInlineMessage(roleEvent, preview(payload, 88, 4)) + a.appendActivity("tool", "Tool output", preview(payload, 88, 4), false) } case agentruntime.EventAgentDone: a.state.IsAgentRunning = false @@ -479,6 +489,7 @@ func (a *App) handleRuntimeEvent(event agentruntime.RuntimeEvent) { } if payload, ok := event.Payload.(provider.Message); ok && strings.TrimSpace(payload.Content) != "" && !a.lastAssistantMatches(payload.Content) { a.activeMessages = append(a.activeMessages, provider.Message{Role: roleAssistant, Content: payload.Content}) + transcriptDirty = true } case agentruntime.EventRunCanceled: a.state.IsAgentRunning = false @@ -486,7 +497,7 @@ func (a *App) handleRuntimeEvent(event agentruntime.RuntimeEvent) { a.state.CurrentTool = "" a.state.ExecutionError = "" a.state.StatusText = statusCanceled - a.appendInlineMessage(roleSystem, "Canceled current run.") + a.appendActivity("run", "Canceled current run", "", false) case agentruntime.EventError: a.state.StatusText = statusError a.state.IsAgentRunning = false @@ -495,7 +506,12 @@ func (a *App) handleRuntimeEvent(event agentruntime.RuntimeEvent) { if payload, ok := event.Payload.(string); ok { a.state.ExecutionError = payload a.state.StatusText = payload - a.appendInlineMessage(roleError, payload) + a.appendActivity("run", "Runtime error", payload, true) + } + case agentruntime.EventProviderRetry: + if payload, ok := event.Payload.(string); ok && strings.TrimSpace(payload) != "" { + a.state.StatusText = statusThinking + a.appendActivity("provider", "Retrying provider call", payload, false) } case agentruntime.EventCompactDone: payload, ok := event.Payload.(agentruntime.CompactDonePayload) @@ -539,6 +555,8 @@ func (a *App) handleRuntimeEvent(event agentruntime.RuntimeEvent) { a.state.StatusText = message a.appendInlineMessage(roleError, message) } + + return transcriptDirty } func (a *App) appendAssistantChunk(chunk string) { @@ -564,6 +582,29 @@ func (a *App) appendInlineMessage(role string, message string) { a.activeMessages = append(a.activeMessages, provider.Message{Role: role, Content: content}) } +func (a *App) appendActivity(kind string, title string, detail string, isError bool) { + title = strings.TrimSpace(title) + detail = strings.TrimSpace(detail) + if title == "" && detail == "" { + return + } + if title == "" { + title = detail + detail = "" + } + + a.activities = append(a.activities, activityEntry{ + Time: time.Now(), + Kind: strings.TrimSpace(kind), + Title: title, + Detail: detail, + IsError: isError, + }) + if len(a.activities) > maxActivityEntries { + a.activities = append([]activityEntry(nil), a.activities[len(a.activities)-maxActivityEntries:]...) + } +} + func (a *App) lastAssistantMatches(content string) bool { if len(a.activeMessages) == 0 { return false @@ -580,9 +621,9 @@ func (a *App) handleViewportKeys(vp *viewport.Model, msg tea.KeyMsg) { case key.Matches(msg, a.keys.ScrollDown): vp.LineDown(2) case key.Matches(msg, a.keys.PageUp): - vp.HalfViewUp() + vp.ViewUp() case key.Matches(msg, a.keys.PageDown): - vp.HalfViewDown() + vp.ViewDown() case key.Matches(msg, a.keys.Top): vp.GotoTop() case key.Matches(msg, a.keys.Bottom): @@ -590,6 +631,49 @@ func (a *App) handleViewportKeys(vp *viewport.Model, msg tea.KeyMsg) { } } +func (a *App) handleTranscriptMouse(msg tea.MouseMsg) bool { + if !a.isMouseWithinTranscript(msg) { + return false + } + + switch msg.Button { + case tea.MouseButtonWheelUp: + a.transcript.LineUp(mouseWheelStepLines) + return true + case tea.MouseButtonWheelDown: + a.transcript.LineDown(mouseWheelStepLines) + return true + default: + return false + } +} + +func (a App) isMouseWithinTranscript(msg tea.MouseMsg) bool { + x, y, width, height := a.transcriptBounds() + if width <= 0 || height <= 0 { + return false + } + return msg.X >= x && msg.X < x+width && msg.Y >= y && msg.Y < y+height +} + +func (a App) transcriptBounds() (int, int, int, int) { + lay := a.computeLayout() + contentX := a.styles.doc.GetPaddingLeft() + contentY := a.styles.doc.GetPaddingTop() + headerHeight := lipgloss.Height(a.renderHeader(lay.contentWidth)) + bodyY := contentY + headerHeight + + streamX := contentX + streamY := bodyY + if lay.stacked { + streamY += lay.sidebarHeight + } else { + streamX += lay.sidebarWidth + lay.bodyGap + } + + return streamX, streamY, lay.rightWidth, a.transcript.Height +} + func (a *App) focusNext() { order := []panel{panelSessions, panelTranscript, panelInput} current := 0 @@ -650,11 +734,12 @@ func (a *App) applyComponentLayout(rebuildTranscript bool) { sidebarBodyHeight := max(4, lay.sidebarHeight-sidebarFrameHeight-lipgloss.Height(a.renderSidebarHeader(sidebarBodyWidth))) a.sessions.SetSize(sidebarBodyWidth, sidebarBodyHeight) menuHeight := a.commandMenuHeight(max(24, lay.rightWidth)) + activityHeight := a.activityPreviewHeight() a.transcript.Width = max(24, lay.rightWidth) a.input.SetWidth(a.composerInnerWidth(lay.rightWidth)) a.input.SetHeight(a.composerHeight()) promptHeight := lipgloss.Height(a.renderPrompt(a.transcript.Width)) - a.transcript.Height = max(6, lay.rightHeight-menuHeight-promptHeight) + a.transcript.Height = max(6, lay.rightHeight-activityHeight-menuHeight-promptHeight) a.providerPicker.SetSize(max(24, clamp(lay.rightWidth-14, 28, 52)), max(4, clamp(lay.rightHeight-10, 6, 10))) a.modelPicker.SetSize(max(24, clamp(lay.rightWidth-14, 28, 52)), max(4, clamp(lay.rightHeight-10, 6, 10))) if rebuildTranscript || prevTranscriptWidth != a.transcript.Width { @@ -766,6 +851,7 @@ func (a *App) startDraftSession() { a.state.ActiveSessionID = "" a.state.ActiveSessionTitle = draftSessionTitle a.activeMessages = nil + a.activities = nil a.state.StatusText = statusDraft a.state.ExecutionError = "" a.state.CurrentTool = "" diff --git a/internal/tui/update_test.go b/internal/tui/update_test.go index a1e649ed..e5497669 100644 --- a/internal/tui/update_test.go +++ b/internal/tui/update_test.go @@ -250,6 +250,12 @@ func TestAppUpdateModelPickerAndRuntimeMessages(t *testing.T) { if app.state.ExecutionError != "" || app.state.StatusText != statusCanceled { t.Fatalf("expected canceled status without error, got %+v", app.state) } + if len(app.activeMessages) != 0 { + t.Fatalf("expected cancel notice to stay out of transcript, got %+v", app.activeMessages) + } + if len(app.activities) == 0 || app.activities[len(app.activities)-1].Title != "Canceled current run" { + t.Fatalf("expected cancel notice in activity, got %+v", app.activities) + } }, }, { @@ -274,6 +280,12 @@ func TestAppUpdateModelPickerAndRuntimeMessages(t *testing.T) { if app.state.StatusText != statusToolError { t.Fatalf("expected status tool error, got %q", app.state.StatusText) } + if len(app.activeMessages) == 0 || app.activeMessages[len(app.activeMessages)-1].Role != roleTool { + t.Fatalf("expected tool result to stay in transcript, got %+v", app.activeMessages) + } + if len(app.activities) == 0 || app.activities[len(app.activities)-1].Title != "Tool error" { + t.Fatalf("expected tool error in activity, got %+v", app.activities) + } }, }, } @@ -1094,6 +1106,12 @@ func TestAppHandleRuntimeEventAdditionalBranches(t *testing.T) { if app.state.CurrentTool != "filesystem_edit" || app.state.StatusText != statusRunningTool { t.Fatalf("unexpected tool start state: %+v", app.state) } + if len(app.activeMessages) != 0 { + t.Fatalf("expected tool start to stay out of transcript, got %+v", app.activeMessages) + } + if len(app.activities) == 0 || app.activities[len(app.activities)-1].Title != "Running tool" { + t.Fatalf("expected tool start in activity, got %+v", app.activities) + } }, }, { @@ -1113,6 +1131,12 @@ func TestAppHandleRuntimeEventAdditionalBranches(t *testing.T) { if app.state.CurrentTool != "" || app.state.StatusText != statusToolFinished { t.Fatalf("unexpected tool success state: %+v", app.state) } + if len(app.activeMessages) == 0 || app.activeMessages[len(app.activeMessages)-1].Role != roleTool { + t.Fatalf("expected tool result message in transcript, got %+v", app.activeMessages) + } + if len(app.activities) == 0 || app.activities[len(app.activities)-1].Title != "Completed tool" { + t.Fatalf("expected tool completion in activity, got %+v", app.activities) + } }, }, { @@ -1127,6 +1151,46 @@ func TestAppHandleRuntimeEventAdditionalBranches(t *testing.T) { if app.state.ExecutionError != "boom" || app.state.IsAgentRunning { t.Fatalf("unexpected error state: %+v", app.state) } + if len(app.activeMessages) != 0 { + t.Fatalf("expected runtime error to stay out of transcript, got %+v", app.activeMessages) + } + if len(app.activities) == 0 || app.activities[len(app.activities)-1].Title != "Runtime error" { + t.Fatalf("expected runtime error in activity, got %+v", app.activities) + } + }, + }, + { + name: "tool call thinking is tracked as activity", + event: agentruntime.RuntimeEvent{ + Type: agentruntime.EventToolCallThinking, + SessionID: "s1", + Payload: "filesystem_edit", + }, + assert: func(t *testing.T, app App) { + t.Helper() + if app.state.CurrentTool != "filesystem_edit" { + t.Fatalf("expected current tool to be populated, got %+v", app.state) + } + if len(app.activities) == 0 || app.activities[len(app.activities)-1].Title != "Planning tool call" { + t.Fatalf("expected planning activity, got %+v", app.activities) + } + }, + }, + { + name: "provider retry is tracked as activity", + event: agentruntime.RuntimeEvent{ + Type: agentruntime.EventProviderRetry, + SessionID: "s1", + Payload: "retrying provider call (attempt 1/2, wait=1.0s)...", + }, + assert: func(t *testing.T, app App) { + t.Helper() + if app.state.StatusText != statusThinking { + t.Fatalf("expected provider retry to preserve thinking status, got %+v", app.state) + } + if len(app.activities) == 0 || app.activities[len(app.activities)-1].Title != "Retrying provider call" { + t.Fatalf("expected provider retry activity, got %+v", app.activities) + } }, }, { @@ -1384,8 +1448,11 @@ func TestAdditionalRenderingAndToolChunkBranches(t *testing.T) { if app.state.StatusText != statusRunningTool { t.Fatalf("expected tool chunk to keep running status, got %q", app.state.StatusText) } - if len(app.activeMessages) == 0 || !strings.Contains(app.activeMessages[len(app.activeMessages)-1].Content, "chunk output") { - t.Fatalf("expected tool chunk preview to append inline message") + if len(app.activeMessages) != 0 { + t.Fatalf("expected tool chunk to stay out of transcript, got %+v", app.activeMessages) + } + if len(app.activities) == 0 || !strings.Contains(app.activities[len(app.activities)-1].Detail, "chunk output") { + t.Fatalf("expected tool chunk preview in activity, got %+v", app.activities) } if got := wrapCodeBlock("a\tb", 3); !strings.Contains(got, "\n") { @@ -1401,6 +1468,81 @@ func TestAdditionalRenderingAndToolChunkBranches(t *testing.T) { } } +func TestHandleViewportKeysPageScrollingUsesFullPage(t *testing.T) { + manager := newTestConfigManager(t) + runtime := newStubRuntime() + app, err := New(nil, manager, runtime, newTestProviderService(t, manager)) + if err != nil { + t.Fatalf("New() error = %v", err) + } + + app.transcript.SetContent(strings.Repeat("line\n", 120)) + app.transcript.Height = 10 + app.transcript.GotoTop() + + app.handleViewportKeys(&app.transcript, tea.KeyMsg{Type: tea.KeyPgDown}) + if app.transcript.YOffset != 10 { + t.Fatalf("expected page down to move a full page, got offset %d", app.transcript.YOffset) + } + + app.handleViewportKeys(&app.transcript, tea.KeyMsg{Type: tea.KeyPgUp}) + if app.transcript.YOffset != 0 { + t.Fatalf("expected page up to return a full page, got offset %d", app.transcript.YOffset) + } +} + +func TestTranscriptMouseWheelScrollsOnlyInsideTranscript(t *testing.T) { + manager := newTestConfigManager(t) + runtime := newStubRuntime() + app, err := New(nil, manager, runtime, newTestProviderService(t, manager)) + if err != nil { + t.Fatalf("New() error = %v", err) + } + + app.width = 128 + app.height = 40 + app.resizeComponents() + app.transcript.SetContent(strings.Repeat("line\n", 160)) + app.transcript.GotoTop() + + x, y, _, _ := app.transcriptBounds() + model, cmd := app.Update(tea.MouseMsg{ + X: x + 1, + Y: y + 1, + Button: tea.MouseButtonWheelDown, + Type: tea.MouseWheelDown, + }) + app = model.(App) + _ = collectTeaMessages(cmd) + if app.transcript.YOffset != mouseWheelStepLines { + t.Fatalf("expected wheel down to scroll transcript by %d lines, got %d", mouseWheelStepLines, app.transcript.YOffset) + } + + model, cmd = app.Update(tea.MouseMsg{ + X: x + 1, + Y: y + 1, + Button: tea.MouseButtonWheelUp, + Type: tea.MouseWheelUp, + }) + app = model.(App) + _ = collectTeaMessages(cmd) + if app.transcript.YOffset != 0 { + t.Fatalf("expected wheel up to scroll transcript back to top, got %d", app.transcript.YOffset) + } + + model, cmd = app.Update(tea.MouseMsg{ + X: 0, + Y: 0, + Button: tea.MouseButtonWheelDown, + Type: tea.MouseWheelDown, + }) + app = model.(App) + _ = collectTeaMessages(cmd) + if app.transcript.YOffset != 0 { + t.Fatalf("expected wheel event outside transcript to be ignored, got %d", app.transcript.YOffset) + } +} + func TestWorkspaceCommandAndFileReferenceFlow(t *testing.T) { previousExecutor := workspaceCommandExecutor t.Cleanup(func() { workspaceCommandExecutor = previousExecutor }) @@ -1431,16 +1573,19 @@ func TestWorkspaceCommandAndFileReferenceFlow(t *testing.T) { if app.state.StatusText != statusCommandDone { t.Fatalf("expected command done status, got %q", app.state.StatusText) } - if len(app.activeMessages) < 2 { - t.Fatalf("expected running event and command result, got %+v", app.activeMessages) + if len(app.activeMessages) != 0 { + t.Fatalf("expected workspace command flow to stay out of transcript, got %+v", app.activeMessages) + } + if len(app.activities) < 2 { + t.Fatalf("expected running event and command result in activity, got %+v", app.activities) } - first := app.activeMessages[0] - if first.Role != roleEvent || !strings.Contains(first.Content, "Running command: git status") { - t.Fatalf("expected running command notice, got %+v", first) + first := app.activities[0] + if first.Title != "Running command" || !strings.Contains(first.Detail, "git status") { + t.Fatalf("expected running command activity, got %+v", first) } - last := app.activeMessages[len(app.activeMessages)-1] - if last.Role != roleTool || !strings.Contains(last.Content, "Command: & git status") || !strings.Contains(last.Content, "stubbed output for git status") { - t.Fatalf("expected tool transcript entry with command output, got %+v", last) + last := app.activities[len(app.activities)-1] + if last.Title != "Command finished" || !strings.Contains(last.Detail, "Command: & git status") || !strings.Contains(last.Detail, "stubbed output for git status") { + t.Fatalf("expected command output in activity, got %+v", last) } app.fileCandidates = []string{"README.md", "internal/tui/update.go", "internal/tui/view.go"} diff --git a/internal/tui/view.go b/internal/tui/view.go index e848fa8f..6b31ae7a 100644 --- a/internal/tui/view.go +++ b/internal/tui/view.go @@ -130,6 +130,9 @@ func (a App) renderWaterfall(width int, height int) string { transcript := a.styles.streamContent.Width(width).Height(a.transcript.Height).Render(a.transcript.View()) parts := []string{transcript} + if activity := a.renderActivityPreview(width); activity != "" { + parts = append(parts, activity) + } if menu := a.renderCommandMenu(width); menu != "" { parts = append(parts, menu) } @@ -418,11 +421,58 @@ func (a App) focusLabel() string { return focusLabelSessions case panelTranscript: return focusLabelTranscript + case panelActivity: + return focusLabelActivity default: return focusLabelComposer } } +func (a App) activityPreviewHeight() int { + if len(a.activities) == 0 { + return 0 + } + return 6 +} + +func (a App) renderActivityPreview(width int) string { + if len(a.activities) == 0 { + return "" + } + + entries := a.activities + if len(entries) > activityPreviewEntries { + entries = entries[len(entries)-activityPreviewEntries:] + } + + lines := make([]string, 0, len(entries)) + bodyWidth := max(10, width-4) + for _, entry := range entries { + lines = append(lines, a.renderActivityLine(entry, bodyWidth)) + } + + return a.renderPanel( + activityTitle, + activitySubtitle, + strings.Join(lines, "\n"), + width, + a.activityPreviewHeight(), + a.focus == panelActivity, + ) +} + +func (a App) renderActivityLine(entry activityEntry, width int) string { + timeLabel := entry.Time.Format("15:04:05") + kindLabel := strings.ToUpper(fallback(strings.TrimSpace(entry.Kind), "event")) + + text := entry.Title + if strings.TrimSpace(entry.Detail) != "" { + text = text + ": " + entry.Detail + } + + return trimMiddle(timeLabel+" "+kindLabel+" "+strings.Join(strings.Fields(text), " "), max(12, width)) +} + func (a App) computeLayout() layout { contentWidth := max(0, a.width-a.styles.doc.GetHorizontalFrameSize()) headerHeight := lipgloss.Height(a.renderHeader(contentWidth)) From fc11903213d0cac4d0d22b492675a27ff68a8ae8 Mon Sep 17 00:00:00 2001 From: creatang Date: Thu, 2 Apr 2026 09:45:17 +0800 Subject: [PATCH 14/23] test(tui): cover activity preview and mouse scroll --- internal/tui/update_test.go | 85 +++++++++++++++++++++++++++++++++++++ 1 file changed, 85 insertions(+) diff --git a/internal/tui/update_test.go b/internal/tui/update_test.go index e5497669..5191f621 100644 --- a/internal/tui/update_test.go +++ b/internal/tui/update_test.go @@ -1489,6 +1489,16 @@ func TestHandleViewportKeysPageScrollingUsesFullPage(t *testing.T) { if app.transcript.YOffset != 0 { t.Fatalf("expected page up to return a full page, got offset %d", app.transcript.YOffset) } + + app.handleViewportKeys(&app.transcript, tea.KeyMsg{Type: tea.KeyEnd}) + if !app.transcript.AtBottom() { + t.Fatalf("expected end to jump to bottom") + } + + app.handleViewportKeys(&app.transcript, tea.KeyMsg{Type: tea.KeyHome}) + if !app.transcript.AtTop() { + t.Fatalf("expected home to jump to top") + } } func TestTranscriptMouseWheelScrollsOnlyInsideTranscript(t *testing.T) { @@ -1541,6 +1551,81 @@ func TestTranscriptMouseWheelScrollsOnlyInsideTranscript(t *testing.T) { if app.transcript.YOffset != 0 { t.Fatalf("expected wheel event outside transcript to be ignored, got %d", app.transcript.YOffset) } + + app.transcript.Height = 0 + if app.isMouseWithinTranscript(tea.MouseMsg{X: x + 1, Y: y + 1}) { + t.Fatalf("expected zero-height transcript bounds to reject mouse hits") + } + + app.transcript.Height = 10 + if app.handleTranscriptMouse(tea.MouseMsg{X: x + 1, Y: y + 1, Button: tea.MouseButtonLeft}) { + t.Fatalf("expected non-wheel mouse button to be ignored") + } + + app.width = 100 + app.height = 32 + app.resizeComponents() + stackX, stackY, _, stackH := app.transcriptBounds() + if stackH <= 0 || !app.isMouseWithinTranscript(tea.MouseMsg{X: stackX + 1, Y: stackY + 1}) { + t.Fatalf("expected stacked layout transcript bounds to accept mouse hits") + } +} + +func TestViewActivityPreviewAndStatusHelpers(t *testing.T) { + manager := newTestConfigManager(t) + runtime := newStubRuntime() + app, err := New(nil, manager, runtime, newTestProviderService(t, manager)) + if err != nil { + t.Fatalf("New() error = %v", err) + } + + if app.activityPreviewHeight() != 0 || app.renderActivityPreview(80) != "" { + t.Fatalf("expected empty activity state to render nothing") + } + + fixed := time.Date(2026, 4, 2, 9, 30, 0, 0, time.UTC) + app.activities = []activityEntry{ + {Time: fixed, Kind: "tool", Title: "first", Detail: "alpha"}, + {Time: fixed, Kind: "", Title: "second", Detail: ""}, + {Time: fixed, Kind: "provider", Title: "third", Detail: "retry"}, + {Time: fixed, Kind: "run", Title: "fourth", Detail: "done"}, + } + app.focus = panelActivity + if app.focusLabel() != focusLabelActivity { + t.Fatalf("expected activity focus label, got %q", app.focusLabel()) + } + if app.activityPreviewHeight() != 6 { + t.Fatalf("expected fixed activity preview height, got %d", app.activityPreviewHeight()) + } + + preview := app.renderActivityPreview(64) + if !strings.Contains(preview, "second") || !strings.Contains(preview, "third") || !strings.Contains(preview, "fourth") { + t.Fatalf("expected last activity entries in preview, got %q", preview) + } + if strings.Contains(preview, "first") { + t.Fatalf("expected oldest activity entry to be trimmed from preview, got %q", preview) + } + + line := app.renderActivityLine(activityEntry{Time: fixed, Kind: "", Title: "single line", Detail: ""}, 80) + if !strings.Contains(line, "EVENT") || strings.Contains(line, "single line:") { + t.Fatalf("expected fallback kind without detail suffix, got %q", line) + } + + rendered := app.renderMessageContent("before\n```go\nfmt.Println(1)\n```\nafter", 30, app.styles.messageBody) + if !strings.Contains(rendered, "before") || !strings.Contains(rendered, "fmt.Println(1)") || !strings.Contains(rendered, "after") { + t.Fatalf("expected mixed prose and code to render, got %q", rendered) + } + + if got := compactStatusText("\n hello world \n", 0); got != "hello world" { + t.Fatalf("expected compact status without truncation, got %q", got) + } + if got := compactStatusText("\n \n", 10); got != "" { + t.Fatalf("expected empty compact status for blank input, got %q", got) + } + + if app.statusBadge("failed request") == "" || app.statusBadge("canceled") == "" || app.statusBadge("ready") == "" { + t.Fatalf("expected status badge branches to render non-empty output") + } } func TestWorkspaceCommandAndFileReferenceFlow(t *testing.T) { From 6a64baa895519e76fe39934bd3eb061038834cfd Mon Sep 17 00:00:00 2001 From: phantom5099 <1011668688@qq.com> Date: Thu, 2 Apr 2026 19:26:20 +0800 Subject: [PATCH 15/23] =?UTF-8?q?refactor(provider):=E5=88=A0=E5=8E=BB?= =?UTF-8?q?=E6=8E=A5=E6=94=B6=E7=9A=84=E6=97=A0=E7=94=A8provider=E5=85=83?= =?UTF-8?q?=E6=95=B0=E6=8D=AE?= MIME-Version: 1.0 Content-Type: text/plain; charset=UTF-8 Content-Transfer-Encoding: 8bit --- docs/guides/adding-providers.md | 24 +-- docs/guides/configuration.md | 11 +- internal/app/bootstrap.go | 16 +- internal/app/bootstrap_test.go | 46 +++++ internal/config/builtin_providers_test.go | 32 +-- internal/config/config_test.go | 80 ++++---- internal/config/loader_test.go | 85 +++++++- internal/config/model.go | 153 ++++---------- internal/provider/metadata.go | 59 +----- internal/provider/metadata_test.go | 19 +- internal/provider/model_catalog_store_test.go | 19 +- internal/provider/openai/openai.go | 2 +- internal/provider/registry_test.go | 158 +++++++------- internal/provider/service.go | 79 +++++-- internal/provider/service_test.go | 192 ++++++++++-------- internal/provider/types.go | 1 - internal/tui/update_test.go | 65 ++++-- 17 files changed, 542 insertions(+), 499 deletions(-) diff --git a/docs/guides/adding-providers.md b/docs/guides/adding-providers.md index 8880a882..db3a0995 100644 --- a/docs/guides/adding-providers.md +++ b/docs/guides/adding-providers.md @@ -10,12 +10,12 @@ NeoCode 的 provider 架构采用**集中式配置 + 驱动复用**的设计： - **配置层**（`internal/config/builtin_providers.go`）：集中管理所有 provider 的元数据 - Provider 名称、Driver 名称 - - Base URL、默认模型、模型列表 - - API Key 环境变量名 + - Base URL、默认模型、API Key 环境变量名 - **驱动层**（`internal/provider/openai/` 等）：负责实际的 API 协议实现 - 请求构造、响应解析 - 流式输出、Tool Call 处理 + - 动态模型发现（如 `GET /models`）与结果归一化 ### 驱动复用 @@ -46,11 +46,6 @@ const ( DeepSeekDefaultAPIKeyEnv = "DEEPSEEK_API_KEY" ) -var deepSeekModels = []string{ - DeepSeekDefaultModel, - "deepseek-coder", -} - // DeepSeekProvider 返回 DeepSeek provider 的默认配置。 func DeepSeekProvider() ProviderConfig { return ProviderConfig{ @@ -58,7 +53,6 @@ func DeepSeekProvider() ProviderConfig { Driver: "openai", // 复用 openai 驱动 BaseURL: DeepSeekDefaultBaseURL, Model: DeepSeekDefaultModel, - Models: append([]string(nil), deepSeekModels...), APIKeyEnv: DeepSeekDefaultAPIKeyEnv, } } @@ -182,19 +176,12 @@ const ( AnthropicDefaultAPIKeyEnv = "ANTHROPIC_API_KEY" ) -var anthropicModels = []string{ - AnthropicDefaultModel, - "claude-opus-4-20250514", - "claude-3-5-sonnet-20241022", -} - func AnthropicProvider() ProviderConfig { return ProviderConfig{ Name: AnthropicName, Driver: "anthropic", // 使用新的 anthropic 驱动 BaseURL: AnthropicDefaultBaseURL, Model: AnthropicDefaultModel, - Models: append([]string(nil), anthropicModels...), APIKeyEnv: AnthropicDefaultAPIKeyEnv, } } @@ -226,7 +213,7 @@ func DefaultProviders() []ProviderConfig { | `ChatRequest` | `internal/provider/types.go` | 请求：`Model`、`SystemPrompt`、`Messages`、`Tools` | | `ChatResponse` | `internal/provider/types.go` | 响应：`Message`、`FinishReason`、`Usage` | | `StreamEvent` | `internal/provider/types.go` | 流式事件：`TextDelta`、`ToolCallStart` | -| `ProviderConfig` | `internal/config/model.go` | 配置：`Name`、`Driver`、`BaseURL`、`Model`、`Models`、`APIKeyEnv` | +| `ProviderConfig` | `internal/config/model.go` | 配置：`Name`、`Driver`、`BaseURL`、`Model`、`APIKeyEnv` | ## 设计约束 @@ -242,7 +229,7 @@ func DefaultProviders() []ProviderConfig { ✅ **驱动职责清晰** - 驱动只负责协议构造与响应解析 -- 不持有模型列表、base_url 等配置信息 +- 不持有 provider 元数据；模型目录由 driver 发现，缓存与合并由 service 处理 ✅ **架构分层** - 厂商差异收敛在 `internal/provider/` 内 @@ -274,9 +261,10 @@ func DefaultProviders() []ProviderConfig { OpenAIProvider(), // OpenAI 官方 API GeminiProvider(), // Google Gemini (OpenAI-compatible) OpenLLProvider(), // OpenLL 服务 (OpenAI-compatible) + QiniuProvider(), // 七牛云推理服务 (OpenAI-compatible) } } ``` -所有内置 provider 都复用 `openai` 驱动，配置集中在 `builtin_providers.go`。 +所有内置 provider 都通过代码集中注册。模型选择器展示的候选模型由默认模型、动态发现结果和本地缓存共同组成。 diff --git a/docs/guides/configuration.md b/docs/guides/configuration.md index 226b9e5b..b668c8be 100644 --- a/docs/guides/configuration.md +++ b/docs/guides/configuration.md @@ -16,6 +16,7 @@ NeoCode 采用"内置 Provider 优先"的策略： - `config.yaml` 不再持久化完整 `providers` 列表 - 只保存当前选择状态和通用运行配置 - 运行时的 `providers` 完全来自代码内置定义 +- 用户不能通过 YAML 注入新的 provider ✅ **安全第一** - API Key 只从环境变量读取，永不写入 YAML @@ -28,6 +29,7 @@ NeoCode 采用"内置 Provider 优先"的策略： | `openai` | `openai` | OpenAI 官方 API | | `gemini` | `openai` | Google Gemini (OpenAI-compatible API) | | `openll` | `openai` | OpenLL 服务 (OpenAI-compatible API) | +| `qiniu` | `openai` | 七牛云推理服务 (OpenAI-compatible API) | 所有内置 provider 都复用 `openai` 驱动，支持流式输出和 Tool Call。 @@ -39,6 +41,7 @@ NeoCode 采用"内置 Provider 优先"的策略： - ✅ 未来代码新增 provider 时，用户无需修改 YAML - ✅ 老配置文件中的 `providers` / `provider_overrides` 会在加载时被清理 - ✅ 配置文件始终保持简洁，只包含必要的运行时状态 +- ✅ 如需新增 provider，只能通过代码扩展内建列表完成 ## 配置文件 @@ -106,6 +109,7 @@ tools: | `openai` | `OPENAI_API_KEY` | [OpenAI Platform](https://platform.openai.com/api-keys) | | `gemini` | `GEMINI_API_KEY` | [Google AI Studio](https://aistudio.google.com/apikey) | | `openll` | `AI_API_KEY` | OpenLL 服务提供商 | +| `qiniu` | `QINIU_API_KEY` | 七牛云推理服务 | ### 环境变量管理 @@ -156,6 +160,7 @@ NeoCode 提供以下 slash 命令用于快速切换配置： ❯ openai gemini openll + qiniu ``` ### /model - 模型选择器 @@ -166,6 +171,8 @@ NeoCode 提供以下 slash 命令用于快速切换配置： 打开当前 provider 的模型选择器。 +模型列表会合并当前 provider 的默认模型、服务端动态发现结果和本地缓存结果。 + **界面示例**（选择 openai provider 后）： ``` ? Select a model: @@ -208,9 +215,9 @@ NeoCode 提供以下 slash 命令用于快速切换配置： 2. 立即持久化到 `config.yaml` 3. 下次启动自动恢复选择状态 -## 添加自定义 Provider +## 扩展内建 Provider -如需添加自定义 provider（如企业内部服务），请参考： +如需添加新的内建 provider（包括企业内部服务），请参考： 👉 [adding-providers.md](./adding-providers.md) diff --git a/internal/app/bootstrap.go b/internal/app/bootstrap.go index 1cc077e3..1ccd1462 100644 --- a/internal/app/bootstrap.go +++ b/internal/app/bootstrap.go @@ -22,22 +22,26 @@ import ( func NewProgram(ctx context.Context) (*tea.Program, error) { loader := config.NewLoader("", builtin.DefaultConfig()) manager := config.NewManager(loader) - cfg, err := manager.Load(ctx) - if err != nil { + if _, err := manager.Load(ctx); err != nil { return nil, err } - toolRegistry := buildToolRegistry(cfg) - toolManager, err := buildToolManager(toolRegistry) + providerRegistry, err := builtin.NewRegistry() if err != nil { return nil, err } + providerService := provider.NewService(manager, providerRegistry, nil) + if _, err := providerService.EnsureSelection(ctx); err != nil { + return nil, err + } - providerRegistry, err := builtin.NewRegistry() + cfg := manager.Get() + + toolRegistry := buildToolRegistry(cfg) + toolManager, err := buildToolManager(toolRegistry) if err != nil { return nil, err } - providerService := provider.NewService(manager, providerRegistry, nil) sessionStore := agentruntime.NewSessionStore(loader.BaseDir()) runtimeSvc := agentruntime.NewWithFactory( diff --git a/internal/app/bootstrap_test.go b/internal/app/bootstrap_test.go index 8b2194bb..1743a0a6 100644 --- a/internal/app/bootstrap_test.go +++ b/internal/app/bootstrap_test.go @@ -7,6 +7,7 @@ import ( "net/http/httptest" "os" "path/filepath" + "strings" "testing" "neo-code/internal/config" @@ -14,6 +15,8 @@ import ( ) func TestNewProgram(t *testing.T) { + disableBuiltinProviderAPIKeys(t) + home := t.TempDir() t.Setenv("HOME", home) t.Setenv("USERPROFILE", home) @@ -32,6 +35,41 @@ func TestNewProgram(t *testing.T) { } } +func TestNewProgramNormalizesInvalidCurrentModelOnStartup(t *testing.T) { + disableBuiltinProviderAPIKeys(t) + + home := t.TempDir() + t.Setenv("HOME", home) + t.Setenv("USERPROFILE", home) + + configDir := filepath.Join(home, ".neocode") + if err := os.MkdirAll(configDir, 0o755); err != nil { + t.Fatalf("mkdir config dir: %v", err) + } + + configPath := filepath.Join(configDir, "config.yaml") + raw := []byte("selected_provider: openai\ncurrent_model: unsupported-current\nworkdir: .\nshell: powershell\n") + if err := os.WriteFile(configPath, raw, 0o644); err != nil { + t.Fatalf("write config: %v", err) + } + + program, err := NewProgram(context.Background()) + if err != nil { + t.Fatalf("NewProgram() error = %v", err) + } + if program == nil { + t.Fatalf("expected tea program") + } + + data, err := os.ReadFile(configPath) + if err != nil { + t.Fatalf("read config: %v", err) + } + if !strings.Contains(string(data), "current_model: "+config.OpenAIDefaultModel) { + t.Fatalf("expected startup normalization to rewrite current_model, got:\n%s", string(data)) + } +} + func TestBuildToolRegistryUsesWebFetchConfig(t *testing.T) { t.Parallel() @@ -126,3 +164,11 @@ func (s stubToolForBootstrap) Schema() map[string]any { return map[string]any{"t func (s stubToolForBootstrap) Execute(ctx context.Context, call tools.ToolCallInput) (tools.ToolResult, error) { return tools.ToolResult{Name: s.name, Content: s.content}, nil } + +func disableBuiltinProviderAPIKeys(t *testing.T) { + t.Helper() + t.Setenv(config.OpenAIDefaultAPIKeyEnv, "") + t.Setenv(config.GeminiDefaultAPIKeyEnv, "") + t.Setenv(config.OpenLLDefaultAPIKeyEnv, "") + t.Setenv(config.QiniuDefaultAPIKeyEnv, "") +} diff --git a/internal/config/builtin_providers_test.go b/internal/config/builtin_providers_test.go index 8e39f9a9..6309a671 100644 --- a/internal/config/builtin_providers_test.go +++ b/internal/config/builtin_providers_test.go @@ -40,12 +40,6 @@ func TestOpenAIProviderConfig(t *testing.T) { if provider.APIKeyEnv != OpenAIDefaultAPIKeyEnv { t.Fatalf("expected API key env %q, got %q", OpenAIDefaultAPIKeyEnv, provider.APIKeyEnv) } - if len(provider.Models) != 0 { - t.Fatalf("expected builtin provider to leave manual models empty, got %+v", provider.Models) - } - if !ContainsModelID(provider.SupportedModels(), OpenAIDefaultModel) { - t.Fatalf("expected supported models to contain default model %q", OpenAIDefaultModel) - } } func TestGeminiProviderConfig(t *testing.T) { @@ -68,12 +62,6 @@ func TestGeminiProviderConfig(t *testing.T) { if provider.APIKeyEnv != GeminiDefaultAPIKeyEnv { t.Fatalf("expected API key env %q, got %q", GeminiDefaultAPIKeyEnv, provider.APIKeyEnv) } - if len(provider.Models) != 0 { - t.Fatalf("expected builtin provider to leave manual models empty, got %+v", provider.Models) - } - if !ContainsModelID(provider.SupportedModels(), GeminiDefaultModel) { - t.Fatalf("expected supported models to contain default model %q", GeminiDefaultModel) - } } func TestOpenLLProviderConfig(t *testing.T) { @@ -96,28 +84,20 @@ func TestOpenLLProviderConfig(t *testing.T) { if provider.APIKeyEnv != OpenLLDefaultAPIKeyEnv { t.Fatalf("expected API key env %q, got %q", OpenLLDefaultAPIKeyEnv, provider.APIKeyEnv) } - if len(provider.Models) != 0 { - t.Fatalf("expected builtin provider to leave manual models empty, got %+v", provider.Models) - } - if !ContainsModelID(provider.SupportedModels(), OpenLLDefaultModel) { - t.Fatalf("expected supported models to contain default model %q", OpenLLDefaultModel) - } } -func TestProviderModelsAreImmutable(t *testing.T) { +func TestProviderDefaultsAreIndependent(t *testing.T) { t.Parallel() provider1 := OpenAIProvider() - models1 := provider1.SupportedModels() - models1[0] = "modified-model" + provider1.Model = "modified-model" provider2 := OpenAIProvider() - models2 := provider2.SupportedModels() - if models2[0] == "modified-model" { - t.Fatal("expected supported models slice to be independent between calls") + if provider2.Model == "modified-model" { + t.Fatal("expected builtin provider defaults to be independent between calls") } - if models2[0] != OpenAIDefaultModel { - t.Fatalf("expected first model %q, got %q", OpenAIDefaultModel, models2[0]) + if provider2.Model != OpenAIDefaultModel { + t.Fatalf("expected default model %q, got %q", OpenAIDefaultModel, provider2.Model) } } diff --git a/internal/config/config_test.go b/internal/config/config_test.go index 1d6b2937..6e515365 100644 --- a/internal/config/config_test.go +++ b/internal/config/config_test.go @@ -22,7 +22,6 @@ func testDefaultProviderConfig() ProviderConfig { Driver: testProviderName, BaseURL: testBaseURL, Model: testModel, - Models: []string{testModel, "gpt-4o", "gpt-5.4", "gpt-5.3-codex"}, APIKeyEnv: testAPIKeyEnv, } } @@ -120,11 +119,11 @@ providers: }, }, { - name: "legacy format", + name: "legacy fields are ignored", data: ` selected_provider: openai current_model: gpt-4o -workspace_root: . +workspace_root: ./definitely-legacy-root shell: bash max_loop: 5 providers: @@ -137,8 +136,8 @@ providers: `, assert: func(t *testing.T, cfg *Config) { t.Helper() - if cfg.MaxLoops != 5 { - t.Fatalf("expected max loops 5, got %d", cfg.MaxLoops) + if cfg.MaxLoops != DefaultMaxLoops { + t.Fatalf("expected legacy max_loop to be ignored, got %d", cfg.MaxLoops) } provider, err := cfg.SelectedProviderConfig() if err != nil { @@ -150,8 +149,8 @@ providers: if cfg.CurrentModel != "gpt-4o" { t.Fatalf("expected current model %q, got %q", "gpt-4o", cfg.CurrentModel) } - if len(provider.Models) != len(testDefaultProviderConfig().Models) { - t.Fatalf("expected builtin provider models, got %+v", provider.Models) + if strings.Contains(filepath.ToSlash(cfg.Workdir), "definitely-legacy-root") { + t.Fatalf("expected legacy workspace_root to be ignored, got %q", cfg.Workdir) } if cfg.Tools.WebFetch.MaxResponseBytes != DefaultWebFetchMaxResponseBytes { t.Fatalf("expected default max_response_bytes %d, got %d", DefaultWebFetchMaxResponseBytes, cfg.Tools.WebFetch.MaxResponseBytes) @@ -167,7 +166,7 @@ providers: tt := tt t.Run(tt.name, func(t *testing.T) { t.Parallel() - cfg, err := parseConfig([]byte(tt.data), *testDefaultConfig()) + cfg, err := parseConfig([]byte(tt.data)) if err != nil { t.Fatalf("parseConfig() error = %v", err) } @@ -296,7 +295,6 @@ func TestManagerConcurrentAccess(t *testing.T) { for k := range next.Providers { if next.Providers[k].Name == next.SelectedProvider { next.Providers[k].Model = model - next.Providers[k].Models = append([]string(nil), models...) } } return nil @@ -509,18 +507,6 @@ func TestProviderConfigValidateFailures(t *testing.T) { }, expectErr: "api_key_env is empty", }, - { - name: "default model missing from models", - provider: ProviderConfig{ - Name: testProviderName, - Driver: testProviderName, - BaseURL: testBaseURL, - Model: testModel, - Models: []string{"gpt-4o"}, - APIKeyEnv: testAPIKeyEnv, - }, - expectErr: "default model \"gpt-4.1\" is not in models", - }, } for _, tt := range tests { @@ -574,7 +560,6 @@ func TestLoaderLoadAndSaveRoundTrip(t *testing.T) { } cfg.CurrentModel = "gpt-5.4" - cfg.Providers[0].Models = []string{"gpt-4.1", "gpt-5.4"} cfg.Providers[0].BaseURL = "https://ignored.example/v1" cfg.Tools.WebFetch.MaxResponseBytes = 1024 cfg.Tools.WebFetch.SupportedContentTypes = []string{"text/html", "application/json"} @@ -614,9 +599,6 @@ func TestLoaderLoadAndSaveRoundTrip(t *testing.T) { if provider.BaseURL != testBaseURL { t.Fatalf("expected provider base url to come from builtin definition, got %q", provider.BaseURL) } - if len(provider.Models) != len(testDefaultProviderConfig().Models) { - t.Fatalf("expected provider models to come from builtin definition, got %+v", provider.Models) - } if reloaded.Tools.WebFetch.MaxResponseBytes != 1024 { t.Fatalf("expected max_response_bytes %d, got %d", 1024, reloaded.Tools.WebFetch.MaxResponseBytes) } @@ -656,25 +638,41 @@ func TestLoaderUsesUpdatedBuiltinProviderWhenUserHasNoOverride(t *testing.T) { } } -func TestApplyProviderDefaultsMatchesByNameOnly(t *testing.T) { - defaults := []ProviderConfig{testDefaultProviderConfig()} +func TestApplyDefaultsReplacesProvidersWithBuiltinSnapshot(t *testing.T) { + t.Parallel() + current := Config{ - Providers: []ProviderConfig{ - { - Name: "openai-alt", - Driver: testProviderName, - }, - }, + Providers: []ProviderConfig{{ + Name: "openai-alt", + Driver: "custom", + BaseURL: "https://example.com/v1", + Model: "custom-model", + APIKeyEnv: "CUSTOM_API_KEY", + }}, + SelectedProvider: "openai-alt", + CurrentModel: "custom-model", } - current.ApplyDefaultsFrom(Config{Providers: defaults}) + current.ApplyDefaultsFrom(*testDefaultConfig()) - provider, err := current.ProviderByName("openai-alt") + if len(current.Providers) != 1 { + t.Fatalf("expected builtin provider snapshot, got %+v", current.Providers) + } + if _, err := current.ProviderByName("openai-alt"); err == nil { + t.Fatalf("expected custom provider to be dropped, got %+v", current.Providers) + } + provider, err := current.ProviderByName(testProviderName) if err != nil { - t.Fatalf("ProviderByName(openai-alt) error = %v", err) + t.Fatalf("ProviderByName(%s) error = %v", testProviderName, err) + } + if provider.BaseURL != testBaseURL || provider.Model != testModel || provider.APIKeyEnv != testAPIKeyEnv { + t.Fatalf("expected builtin provider metadata, got %+v", provider) + } + if current.SelectedProvider != testProviderName { + t.Fatalf("expected selected provider to reset to builtin %q, got %q", testProviderName, current.SelectedProvider) } - if provider.BaseURL != "" || provider.Model != "" || provider.APIKeyEnv != "" || len(provider.Models) != 0 { - t.Fatalf("expected provider metadata to stay empty when only driver matches, got %+v", provider) + if current.CurrentModel != testModel { + t.Fatalf("expected current model to reset with selected builtin provider, got %q", current.CurrentModel) } } @@ -757,13 +755,13 @@ func TestNormalizeWorkdirAndClone(t *testing.T) { cfg := testDefaultConfig() cloned := cfg.Clone() cloned.CurrentModel = "modified" - cloned.Providers[0].Models[0] = "modified-model" + cloned.Providers[0].BaseURL = "https://modified.example/v1" cloned.Tools.WebFetch.SupportedContentTypes[0] = "application/json" if cfg.CurrentModel == cloned.CurrentModel { t.Fatalf("expected clone to be independent from source") } - if cfg.Providers[0].Models[0] == cloned.Providers[0].Models[0] { - t.Fatalf("expected provider models to be cloned") + if cfg.Providers[0].BaseURL == cloned.Providers[0].BaseURL { + t.Fatalf("expected providers to be cloned") } if cfg.Tools.WebFetch.SupportedContentTypes[0] == cloned.Tools.WebFetch.SupportedContentTypes[0] { t.Fatalf("expected webfetch supported content types to be cloned") diff --git a/internal/config/loader_test.go b/internal/config/loader_test.go index 1b54453a..054e9db5 100644 --- a/internal/config/loader_test.go +++ b/internal/config/loader_test.go @@ -79,9 +79,6 @@ providers: type: openai base_url: https://example.com/v1 model: gpt-5.4 - models: - - gpt-4.1 - - gpt-5.4 api_key_env: OPENAI_API_KEY ` if err := os.WriteFile(loader.ConfigPath(), []byte(strings.TrimSpace(legacy)+"\n"), 0o644); err != nil { @@ -102,9 +99,6 @@ providers: if cfg.CurrentModel != "gpt-5.4" { t.Fatalf("expected current model to stay %q, got %q", "gpt-5.4", cfg.CurrentModel) } - if len(provider.Models) != len(testDefaultProviderConfig().Models) { - t.Fatalf("expected builtin provider models, got %+v", provider.Models) - } rewritten, err := os.ReadFile(loader.ConfigPath()) if err != nil { @@ -121,3 +115,82 @@ providers: t.Fatalf("expected rewritten config to omit provider metadata, got:\n%s", text) } } + +func TestLoaderRewritesNormalizedSelectionStateOnLoad(t *testing.T) { + t.Parallel() + + loader := NewLoader(t.TempDir(), testDefaultConfig()) + if err := os.MkdirAll(loader.BaseDir(), 0o755); err != nil { + t.Fatalf("mkdir base dir: %v", err) + } + + raw := ` +selected_provider: missing-provider +workdir: . +shell: powershell +` + if err := os.WriteFile(loader.ConfigPath(), []byte(strings.TrimSpace(raw)+"\n"), 0o644); err != nil { + t.Fatalf("write config: %v", err) + } + + cfg, err := loader.Load(context.Background()) + if err != nil { + t.Fatalf("Load() error = %v", err) + } + if cfg.SelectedProvider != testProviderName { + t.Fatalf("expected selected provider %q, got %q", testProviderName, cfg.SelectedProvider) + } + if cfg.CurrentModel != testModel { + t.Fatalf("expected current model %q, got %q", testModel, cfg.CurrentModel) + } + + rewritten, err := os.ReadFile(loader.ConfigPath()) + if err != nil { + t.Fatalf("read rewritten config: %v", err) + } + text := string(rewritten) + if !strings.Contains(text, "selected_provider: "+testProviderName) { + t.Fatalf("expected rewritten config to persist selected provider, got:\n%s", text) + } + if !strings.Contains(text, "current_model: "+testModel) { + t.Fatalf("expected rewritten config to persist current model, got:\n%s", text) + } +} + +func TestLoaderRewritesMissingCurrentModelOnLoad(t *testing.T) { + t.Parallel() + + loader := NewLoader(t.TempDir(), testDefaultConfig()) + if err := os.MkdirAll(loader.BaseDir(), 0o755); err != nil { + t.Fatalf("mkdir base dir: %v", err) + } + + raw := ` +selected_provider: openai +workdir: . +shell: powershell +` + if err := os.WriteFile(loader.ConfigPath(), []byte(strings.TrimSpace(raw)+"\n"), 0o644); err != nil { + t.Fatalf("write config: %v", err) + } + + cfg, err := loader.Load(context.Background()) + if err != nil { + t.Fatalf("Load() error = %v", err) + } + if cfg.SelectedProvider != testProviderName { + t.Fatalf("expected selected provider %q, got %q", testProviderName, cfg.SelectedProvider) + } + if cfg.CurrentModel != testModel { + t.Fatalf("expected current model %q, got %q", testModel, cfg.CurrentModel) + } + + rewritten, err := os.ReadFile(loader.ConfigPath()) + if err != nil { + t.Fatalf("read rewritten config: %v", err) + } + text := string(rewritten) + if !strings.Contains(text, "current_model: "+testModel) { + t.Fatalf("expected rewritten config to persist current model, got:\n%s", text) + } +} diff --git a/internal/config/model.go b/internal/config/model.go index 3bf19fc9..4341e49c 100644 --- a/internal/config/model.go +++ b/internal/config/model.go @@ -48,12 +48,11 @@ type Config struct { } type ProviderConfig struct { - Name string `yaml:"name"` - Driver string `yaml:"driver"` - BaseURL string `yaml:"base_url"` - Model string `yaml:"model"` - Models []string `yaml:"models,omitempty"` - APIKeyEnv string `yaml:"api_key_env"` + Name string `yaml:"name"` + Driver string `yaml:"driver"` + BaseURL string `yaml:"base_url"` + Model string `yaml:"model"` + APIKeyEnv string `yaml:"api_key_env"` } type ResolvedProviderConfig struct { @@ -117,16 +116,23 @@ func (c *Config) ApplyDefaultsFrom(defaults Config) { return } - if len(c.Providers) == 0 { + if len(defaults.Providers) > 0 { c.Providers = cloneProviders(defaults.Providers) - } else { - c.Providers = applyProviderDefaults(c.Providers, defaults.Providers) } - if strings.TrimSpace(c.SelectedProvider) == "" { - c.SelectedProvider = defaults.SelectedProvider + fallbackProvider := defaultSelectedProviderName(c.Providers, defaults.SelectedProvider) + selectedReset := false + switch current := strings.TrimSpace(c.SelectedProvider); { + case current == "": + c.SelectedProvider = fallbackProvider + selectedReset = true + case !containsProviderName(c.Providers, current): + c.SelectedProvider = fallbackProvider + selectedReset = true + default: + c.SelectedProvider = current } - if strings.TrimSpace(c.CurrentModel) == "" { + if strings.TrimSpace(c.CurrentModel) == "" || selectedReset { if selected, err := c.SelectedProviderConfig(); err == nil { c.CurrentModel = selected.Model } else if strings.TrimSpace(defaults.CurrentModel) != "" { @@ -257,14 +263,6 @@ func (p ProviderConfig) Validate() error { if _, err := p.Identity(); err != nil { return fmt.Errorf("provider %q: %w", p.Name, err) } - if models := normalizeModelIDs(p.Models); len(p.Models) > 0 { - if len(models) == 0 { - return fmt.Errorf("provider %q models is empty", p.Name) - } - if !ContainsModelID(models, p.Model) { - return fmt.Errorf("provider %q default model %q is not in models", p.Name, strings.TrimSpace(p.Model)) - } - } return nil } @@ -272,23 +270,6 @@ func (p ProviderConfig) Identity() (ProviderIdentity, error) { return NewProviderIdentity(p.Driver, p.BaseURL) } -func (p ProviderConfig) ConfiguredModels() []string { - return cloneModelIDs(p.Models) -} - -func (p ProviderConfig) SupportedModels() []string { - models := p.ConfiguredModels() - if len(models) > 0 { - return models - } - - model := strings.TrimSpace(p.Model) - if model == "" { - return nil - } - return []string{model} -} - func (p ProviderConfig) ResolveAPIKey() (string, error) { envName := strings.TrimSpace(p.APIKeyEnv) if envName == "" { @@ -315,57 +296,6 @@ func (p ProviderConfig) Resolve() (ResolvedProviderConfig, error) { }, nil } -func applyProviderDefaults(providers []ProviderConfig, defaults []ProviderConfig) []ProviderConfig { - out := make([]ProviderConfig, 0, len(providers)) - for _, provider := range providers { - out = append(out, mergeProviderDefaults(provider, defaults)) - } - return out -} - -func mergeProviderDefaults(provider ProviderConfig, defaults []ProviderConfig) ProviderConfig { - base, ok := matchDefaultProvider(provider, defaults) - if !ok { - return provider - } - - if strings.TrimSpace(provider.Name) == "" { - provider.Name = base.Name - } - if strings.TrimSpace(provider.Driver) == "" { - provider.Driver = base.Driver - } - if strings.TrimSpace(provider.BaseURL) == "" { - provider.BaseURL = base.BaseURL - } - if strings.TrimSpace(provider.Model) == "" { - provider.Model = base.Model - } - if len(provider.Models) == 0 { - provider.Models = cloneModelIDs(base.Models) - } - if strings.TrimSpace(provider.APIKeyEnv) == "" { - provider.APIKeyEnv = base.APIKeyEnv - } - - return provider -} - -func matchDefaultProvider(provider ProviderConfig, defaults []ProviderConfig) (ProviderConfig, bool) { - name := NormalizeProviderName(provider.Name) - if name == "" { - return ProviderConfig{}, false - } - - for _, candidate := range defaults { - if NormalizeProviderName(candidate.Name) == name { - return candidate, true - } - } - - return ProviderConfig{}, false -} - func normalizeWorkdir(workdir string) string { workdir = strings.TrimSpace(workdir) if workdir == "" { @@ -590,50 +520,41 @@ func cloneProviders(providers []ProviderConfig) []ProviderConfig { cloned := make([]ProviderConfig, 0, len(providers)) for _, provider := range providers { - next := provider - next.Models = cloneModelIDs(provider.Models) - cloned = append(cloned, next) + cloned = append(cloned, provider) } return cloned } -func cloneModelIDs(models []string) []string { - normalized := normalizeModelIDs(models) - if len(normalized) == 0 { - return nil +func defaultSelectedProviderName(providers []ProviderConfig, fallback string) string { + if containsProviderName(providers, fallback) { + return strings.TrimSpace(fallback) } - return append([]string(nil), normalized...) + if len(providers) == 0 { + return "" + } + return strings.TrimSpace(providers[0].Name) } -func normalizeModelIDs(models []string) []string { - if len(models) == 0 { - return nil +func containsProviderName(providers []ProviderConfig, name string) bool { + target := NormalizeProviderName(name) + if target == "" { + return false } - - normalized := make([]string, 0, len(models)) - seen := make(map[string]struct{}, len(models)) - for _, model := range models { - id := strings.TrimSpace(model) - if id == "" { - continue - } - key := NormalizeKey(id) - if _, exists := seen[key]; exists { - continue + for _, provider := range providers { + if NormalizeProviderName(provider.Name) == target { + return true } - seen[key] = struct{}{} - normalized = append(normalized, id) } - return normalized + return false } func ContainsModelID(models []string, model string) bool { - target := strings.ToLower(strings.TrimSpace(model)) + target := NormalizeKey(model) if target == "" { return false } - for _, candidate := range normalizeModelIDs(models) { - if strings.EqualFold(candidate, target) { + for _, candidate := range models { + if NormalizeKey(candidate) == target { return true } } diff --git a/internal/provider/metadata.go b/internal/provider/metadata.go index cf04de01..319a87f9 100644 --- a/internal/provider/metadata.go +++ b/internal/provider/metadata.go @@ -6,8 +6,8 @@ import ( "neo-code/internal/config" ) -// DescriptorFromMetadata normalizes a raw provider model object into a ModelDescriptor. -func DescriptorFromMetadata(raw map[string]any) (ModelDescriptor, bool) { +// DescriptorFromRawModel normalizes a raw provider model object into a ModelDescriptor. +func DescriptorFromRawModel(raw map[string]any) (ModelDescriptor, bool) { id := firstNonEmptyString( stringValue(raw["id"]), stringValue(raw["model"]), @@ -24,7 +24,6 @@ func DescriptorFromMetadata(raw map[string]any) (ModelDescriptor, bool) { ContextWindow: firstPositiveInt(raw["context_window"], raw["contextLength"], raw["input_token_limit"], raw["max_context_tokens"]), MaxOutputTokens: firstPositiveInt(raw["max_output_tokens"], raw["output_token_limit"], raw["max_tokens"]), Capabilities: boolMapValue(raw["capabilities"]), - Metadata: cloneStringAnyMap(raw), } return normalizeModelDescriptor(descriptor), true } @@ -91,7 +90,6 @@ func normalizeModelDescriptor(descriptor ModelDescriptor) ModelDescriptor { descriptor.Name = descriptor.ID } descriptor.Capabilities = cloneStringBoolMap(descriptor.Capabilities) - descriptor.Metadata = cloneStringAnyMap(descriptor.Metadata) return descriptor } @@ -109,7 +107,6 @@ func mergeModelDescriptor(primary ModelDescriptor, secondary ModelDescriptor) Mo primary.MaxOutputTokens = secondary.MaxOutputTokens } primary.Capabilities = mergeStringBoolMaps(primary.Capabilities, secondary.Capabilities) - primary.Metadata = mergeStringAnyMaps(primary.Metadata, secondary.Metadata) return normalizeModelDescriptor(primary) } @@ -131,32 +128,6 @@ func mergeStringBoolMaps(primary map[string]bool, secondary map[string]bool) map return merged } -func mergeStringAnyMaps(primary map[string]any, secondary map[string]any) map[string]any { - if len(primary) == 0 && len(secondary) == 0 { - return nil - } - - merged := cloneStringAnyMap(primary) - if merged == nil { - merged = make(map[string]any, len(secondary)) - } - - for key, value := range secondary { - current, exists := merged[key] - if !exists { - merged[key] = cloneJSONValue(value) - continue - } - - currentMap, currentOK := current.(map[string]any) - valueMap, valueOK := value.(map[string]any) - if currentOK && valueOK { - merged[key] = mergeStringAnyMaps(currentMap, valueMap) - } - } - return merged -} - func boolMapValue(value any) map[string]bool { raw, ok := value.(map[string]any) if !ok || len(raw) == 0 { @@ -234,32 +205,6 @@ func cloneStringBoolMap(source map[string]bool) map[string]bool { return cloned } -func cloneStringAnyMap(source map[string]any) map[string]any { - if len(source) == 0 { - return nil - } - cloned := make(map[string]any, len(source)) - for key, value := range source { - cloned[key] = cloneJSONValue(value) - } - return cloned -} - -func cloneJSONValue(value any) any { - switch typed := value.(type) { - case map[string]any: - return cloneStringAnyMap(typed) - case []any: - cloned := make([]any, 0, len(typed)) - for _, item := range typed { - cloned = append(cloned, cloneJSONValue(item)) - } - return cloned - default: - return typed - } -} - func normalizeModelID(id string) string { return config.NormalizeKey(id) } diff --git a/internal/provider/metadata_test.go b/internal/provider/metadata_test.go index d1033fc1..d205e196 100644 --- a/internal/provider/metadata_test.go +++ b/internal/provider/metadata_test.go @@ -2,7 +2,7 @@ package provider import "testing" -func TestDescriptorFromMetadataPreservesRawFields(t *testing.T) { +func TestDescriptorFromRawModelNormalizesUsefulFields(t *testing.T) { t.Parallel() raw := map[string]any{ @@ -21,7 +21,7 @@ func TestDescriptorFromMetadataPreservesRawFields(t *testing.T) { }, } - descriptor, ok := DescriptorFromMetadata(raw) + descriptor, ok := DescriptorFromRawModel(raw) if !ok { t.Fatal("expected descriptor to be normalized") } @@ -34,12 +34,12 @@ func TestDescriptorFromMetadataPreservesRawFields(t *testing.T) { if !descriptor.Capabilities["tool_call"] { t.Fatalf("expected tool_call capability, got %+v", descriptor.Capabilities) } - if _, ok := descriptor.Metadata["experimental"]; !ok { - t.Fatalf("expected unknown metadata to be preserved, got %+v", descriptor.Metadata) + if _, ok := descriptor.Capabilities["notes"]; ok { + t.Fatalf("expected non-bool capability values to be ignored, got %+v", descriptor.Capabilities) } } -func TestMergeModelDescriptorsPrefersEarlierSourceAndBackfillsMetadata(t *testing.T) { +func TestMergeModelDescriptorsPrefersEarlierSourceAndBackfillsUsefulFields(t *testing.T) { t.Parallel() manual := []ModelDescriptor{{ @@ -51,9 +51,8 @@ func TestMergeModelDescriptorsPrefersEarlierSourceAndBackfillsMetadata(t *testin Name: "Discovered Name", ContextWindow: 64000, MaxOutputTokens: 4096, - Metadata: map[string]any{ - "id": "gpt-test", - "context_window": float64(64000), + Capabilities: map[string]bool{ + "tool_call": true, }, }} @@ -67,7 +66,7 @@ func TestMergeModelDescriptorsPrefersEarlierSourceAndBackfillsMetadata(t *testin if merged[0].ContextWindow != 64000 { t.Fatalf("expected metadata to be backfilled, got %+v", merged[0]) } - if _, ok := merged[0].Metadata["context_window"]; !ok { - t.Fatalf("expected raw metadata to survive merge, got %+v", merged[0].Metadata) + if !merged[0].Capabilities["tool_call"] { + t.Fatalf("expected capabilities to be backfilled, got %+v", merged[0].Capabilities) } } diff --git a/internal/provider/model_catalog_store_test.go b/internal/provider/model_catalog_store_test.go index 73f55396..70b17043 100644 --- a/internal/provider/model_catalog_store_test.go +++ b/internal/provider/model_catalog_store_test.go @@ -24,11 +24,13 @@ func TestJSONModelCatalogStoreRoundTrip(t *testing.T) { FetchedAt: time.Date(2026, 4, 2, 10, 0, 0, 0, time.UTC), ExpiresAt: time.Date(2026, 4, 3, 10, 0, 0, 0, time.UTC), Models: []ModelDescriptor{{ - ID: "gpt-test", - Name: "GPT Test", - Metadata: map[string]any{ - "id": "gpt-test", - "experimental": map[string]any{"tier": "beta"}, + ID: "gpt-test", + Name: "GPT Test", + Description: "normalized model", + ContextWindow: 128000, + MaxOutputTokens: 8192, + Capabilities: map[string]bool{ + "tool_call": true, }, }}, } @@ -47,8 +49,11 @@ func TestJSONModelCatalogStoreRoundTrip(t *testing.T) { if len(got.Models) != 1 { t.Fatalf("expected 1 model, got %+v", got.Models) } - if _, ok := got.Models[0].Metadata["experimental"]; !ok { - t.Fatalf("expected nested metadata to survive round-trip, got %+v", got.Models[0].Metadata) + if got.Models[0].ContextWindow != 128000 || got.Models[0].MaxOutputTokens != 8192 { + t.Fatalf("expected normalized token fields to survive round-trip, got %+v", got.Models[0]) + } + if !got.Models[0].Capabilities["tool_call"] { + t.Fatalf("expected capabilities to survive round-trip, got %+v", got.Models[0].Capabilities) } } diff --git a/internal/provider/openai/openai.go b/internal/provider/openai/openai.go index 0d33348d..c0f891c3 100644 --- a/internal/provider/openai/openai.go +++ b/internal/provider/openai/openai.go @@ -94,7 +94,7 @@ func (p *Provider) DiscoverModels(ctx context.Context) ([]domain.ModelDescriptor descriptors := make([]domain.ModelDescriptor, 0, len(rawModels)) for _, raw := range rawModels { - descriptor, ok := domain.DescriptorFromMetadata(raw) + descriptor, ok := domain.DescriptorFromRawModel(raw) if !ok { continue } diff --git a/internal/provider/registry_test.go b/internal/provider/registry_test.go index 5627e9ca..7f3cb9df 100644 --- a/internal/provider/registry_test.go +++ b/internal/provider/registry_test.go @@ -99,72 +99,47 @@ func TestRegistryRejectsDuplicateDriverRegistration(t *testing.T) { } } -func TestServiceListProvidersUsesConfiguredMetadata(t *testing.T) { +func TestServiceListProvidersFallsBackToProviderDefaultModel(t *testing.T) { t.Parallel() manager := newTestManager(t) - if err := manager.Update(context.Background(), func(cfg *config.Config) error { - cfg.Providers = append(cfg.Providers, config.ProviderConfig{ - Name: "unsupported", - Driver: "custom", - BaseURL: "https://example.com", - Model: "custom-model", - Models: []string{"custom-model"}, - APIKeyEnv: "CUSTOM_API_KEY", - }) - return nil - }); err != nil { - t.Fatalf("append provider: %v", err) - } service := provider.NewService(manager, newTestRegistry(t), newCatalogStoreStub()) items, err := service.ListProviders(context.Background()) if err != nil { t.Fatalf("ListProviders() error = %v", err) } - expectedModels := map[string]int{ - config.OpenAIName: len(config.OpenAIProvider().SupportedModels()), - config.GeminiName: len(config.GeminiProvider().SupportedModels()), - config.OpenLLName: len(config.OpenLLProvider().SupportedModels()), - config.QiniuName: len(config.QiniuProvider().SupportedModels()), + expected := map[string]int{ + config.OpenAIName: 1, + config.GeminiName: 1, + config.OpenLLName: 1, + config.QiniuName: 1, } - if len(items) != len(expectedModels) { - t.Fatalf("expected only supported providers, got %d", len(items)) + if len(items) != len(expected) { + t.Fatalf("expected only builtin providers, got %d", len(items)) } for _, item := range items { - wantModels, ok := expectedModels[item.ID] + wantModels, ok := expected[item.ID] if !ok { - t.Fatalf("unexpected supported provider %q", item.ID) + t.Fatalf("unexpected builtin provider %q", item.ID) } if item.Description != "" { t.Fatalf("expected provider description to stay empty for hidden metadata, got %q", item.Description) } if len(item.Models) != wantModels { - t.Fatalf("expected provider models to come from config, got %+v", item.Models) + t.Fatalf("expected provider models to fall back to the default model, got %+v", item.Models) } - delete(expectedModels, item.ID) + delete(expected, item.ID) } - if len(expectedModels) != 0 { - t.Fatalf("missing supported providers from catalog: %+v", expectedModels) + if len(expected) != 0 { + t.Fatalf("missing builtin providers from catalog: %+v", expected) } } func TestServiceSelectProviderAndSetCurrentModel(t *testing.T) { - defaults := builtin.DefaultConfig() - defaults.Providers = append(defaults.Providers, config.ProviderConfig{ - Name: "custom-main", - Driver: "custom", - BaseURL: "https://example.com", - Model: "custom-model", - Models: []string{"custom-model", "custom-alt"}, - APIKeyEnv: "CUSTOM_API_KEY", - }) - manager := newTestManagerWithDefaults(t, defaults) + manager := newTestManager(t) registry := newTestRegistry(t) - if err := registry.Register(stubDriver("custom")); err != nil { - t.Fatalf("register stub driver: %v", err) - } if err := manager.Update(context.Background(), func(cfg *config.Config) error { cfg.CurrentModel = "gpt-5.4" @@ -173,13 +148,28 @@ func TestServiceSelectProviderAndSetCurrentModel(t *testing.T) { t.Fatalf("seed current model: %v", err) } - service := provider.NewService(manager, registry, newCatalogStoreStub()) + store := newCatalogStoreStub() + identity, err := config.QiniuProvider().Identity() + if err != nil { + t.Fatalf("qiniu provider identity: %v", err) + } + if err := store.Save(context.Background(), provider.ModelCatalog{ + SchemaVersion: 1, + Identity: identity, + Models: []provider.ModelDescriptor{ + {ID: "qiniu-alt", Name: "qiniu-alt"}, + }, + }); err != nil { + t.Fatalf("seed catalog: %v", err) + } + + service := provider.NewService(manager, registry, store) - selection, err := service.SelectProvider(context.Background(), "custom-main") + selection, err := service.SelectProvider(context.Background(), config.QiniuName) if err != nil { t.Fatalf("SelectProvider() error = %v", err) } - if selection.ProviderID != "custom-main" || selection.ModelID != "custom-model" { + if selection.ProviderID != config.QiniuName || selection.ModelID != config.QiniuDefaultModel { t.Fatalf("unexpected selection after switch: %+v", selection) } @@ -187,12 +177,12 @@ func TestServiceSelectProviderAndSetCurrentModel(t *testing.T) { t.Fatalf("expected ErrModelNotFound, got %v", err) } - selection, err = service.SetCurrentModel(context.Background(), "custom-alt") + selection, err = service.SetCurrentModel(context.Background(), "qiniu-alt") if err != nil { t.Fatalf("SetCurrentModel() error = %v", err) } - if selection.ModelID != "custom-alt" { - t.Fatalf("expected selected model %q, got %+v", "custom-alt", selection) + if selection.ModelID != "qiniu-alt" { + t.Fatalf("expected selected model %q, got %+v", "qiniu-alt", selection) } cfg := manager.Get() @@ -200,10 +190,10 @@ func TestServiceSelectProviderAndSetCurrentModel(t *testing.T) { if err != nil { t.Fatalf("SelectedProviderConfig() error = %v", err) } - if selected.Name != "custom-main" { - t.Fatalf("expected selected provider %q, got %+v", "custom-main", selected) + if selected.Name != config.QiniuName { + t.Fatalf("expected selected provider %q, got %+v", config.QiniuName, selected) } - if selected.Model != "custom-model" || cfg.CurrentModel != "custom-alt" { + if selected.Model != config.QiniuDefaultModel || cfg.CurrentModel != "qiniu-alt" { t.Fatalf("expected provider default and current model to diverge safely, got provider=%q current=%q", selected.Model, cfg.CurrentModel) } @@ -215,30 +205,25 @@ func TestServiceSelectProviderAndSetCurrentModel(t *testing.T) { if err != nil { t.Fatalf("SelectedProviderConfig() after reload error = %v", err) } - if selected.Model != "custom-model" || reloaded.CurrentModel != "custom-alt" { + if selected.Model != config.QiniuDefaultModel || reloaded.CurrentModel != "qiniu-alt" { t.Fatalf("expected current model persistence without overriding provider default, got provider=%q current=%q", selected.Model, reloaded.CurrentModel) } } -func TestServiceModelOperationsUseProviderConfigEvenWithoutDriver(t *testing.T) { +func TestServiceModelOperationsUseBuiltinConfigEvenWithoutDriver(t *testing.T) { t.Parallel() - manager := newTestManager(t) - if err := manager.Update(context.Background(), func(cfg *config.Config) error { - cfg.Providers = append(cfg.Providers, config.ProviderConfig{ - Name: "broken-provider", - Driver: "missing-driver", - BaseURL: "https://example.com", - Model: "broken-model", - Models: []string{"broken-model", "broken-alt"}, - APIKeyEnv: "BROKEN_API_KEY", - }) - cfg.SelectedProvider = "broken-provider" - cfg.CurrentModel = "broken-model" - return nil - }); err != nil { - t.Fatalf("append broken provider: %v", err) - } + defaults := builtin.DefaultConfig() + defaults.Providers = []config.ProviderConfig{{ + Name: config.OpenAIName, + Driver: "missing-driver", + BaseURL: config.OpenAIDefaultBaseURL, + Model: "broken-model", + APIKeyEnv: config.OpenAIDefaultAPIKeyEnv, + }} + defaults.SelectedProvider = config.OpenAIName + defaults.CurrentModel = "broken-model" + manager := newTestManagerWithDefaults(t, defaults) service := provider.NewService(manager, newTestRegistry(t), newCatalogStoreStub()) @@ -246,39 +231,32 @@ func TestServiceModelOperationsUseProviderConfigEvenWithoutDriver(t *testing.T) if err != nil { t.Fatalf("ListModels() error = %v", err) } - if len(models) != 2 || models[1].ID != "broken-alt" { - t.Fatalf("expected models from provider config, got %+v", models) + if len(models) != 1 || models[0].ID != "broken-model" { + t.Fatalf("expected default model fallback, got %+v", models) } - selection, err := service.SetCurrentModel(context.Background(), "broken-alt") - if err != nil { - t.Fatalf("SetCurrentModel() error = %v", err) - } - if selection.ModelID != "broken-alt" { - t.Fatalf("expected current model to update from provider config, got %+v", selection) + if _, err := service.SetCurrentModel(context.Background(), "broken-alt"); !errors.Is(err, provider.ErrModelNotFound) { + t.Fatalf("expected SetCurrentModel() to reject undiscovered model, got %v", err) } } func TestServiceSelectProviderRejectsUnsupportedDriver(t *testing.T) { t.Parallel() - manager := newTestManager(t) - if err := manager.Update(context.Background(), func(cfg *config.Config) error { - cfg.Providers = append(cfg.Providers, config.ProviderConfig{ - Name: "broken-provider", - Driver: "missing-driver", - BaseURL: "https://example.com", - Model: "broken-model", - Models: []string{"broken-model"}, - APIKeyEnv: "BROKEN_API_KEY", - }) - return nil - }); err != nil { - t.Fatalf("append broken provider: %v", err) - } + defaults := builtin.DefaultConfig() + defaults.Providers = []config.ProviderConfig{{ + Name: config.OpenAIName, + Driver: "missing-driver", + BaseURL: config.OpenAIDefaultBaseURL, + Model: config.OpenAIDefaultModel, + APIKeyEnv: config.OpenAIDefaultAPIKeyEnv, + }} + defaults.SelectedProvider = config.OpenAIName + defaults.CurrentModel = config.OpenAIDefaultModel + manager := newTestManagerWithDefaults(t, defaults) service := provider.NewService(manager, newTestRegistry(t), newCatalogStoreStub()) - if _, err := service.SelectProvider(context.Background(), "broken-provider"); !errors.Is(err, provider.ErrDriverNotFound) { + if _, err := service.SelectProvider(context.Background(), config.OpenAIName); !errors.Is(err, provider.ErrDriverNotFound) { t.Fatalf("expected SelectProvider() to preserve ErrDriverNotFound, got %v", err) } } diff --git a/internal/provider/service.go b/internal/provider/service.go index e0e8b18e..0dbded09 100644 --- a/internal/provider/service.go +++ b/internal/provider/service.go @@ -94,7 +94,8 @@ func (s *Service) SelectProvider(ctx context.Context, providerName string) (Prov } cfg.SelectedProvider = selected.Name - cfg.CurrentModel = selectModel(cfg.CurrentModel, modelDescriptorIDs(models), selected.Model) + nextModel, _ := resolveCurrentModel(cfg.CurrentModel, models, selected.Model) + cfg.CurrentModel = nextModel selection = selectionFromConfig(*cfg) return nil }) @@ -173,6 +174,45 @@ func (s *Service) Build(ctx context.Context, cfg config.ResolvedProviderConfig) return s.registry.Build(ctx, cfg) } +func (s *Service) EnsureSelection(ctx context.Context) (ProviderSelection, error) { + if err := s.validate(); err != nil { + return ProviderSelection{}, err + } + if err := ctx.Err(); err != nil { + return ProviderSelection{}, err + } + + cfgSnapshot := s.manager.Get() + selected, err := cfgSnapshot.SelectedProviderConfig() + if err != nil { + return ProviderSelection{}, err + } + + models := s.modelsForProvider(ctx, selected, modelQueryOptions{ + allowSyncRefresh: true, + queueRefresh: true, + }) + nextModel, changed := resolveCurrentModel(cfgSnapshot.CurrentModel, models, selected.Model) + if !changed { + return selectionFromConfig(cfgSnapshot), nil + } + + var selection ProviderSelection + err = s.manager.Update(ctx, func(cfg *config.Config) error { + if _, err := cfg.SelectedProviderConfig(); err != nil { + return err + } + cfg.CurrentModel = nextModel + selection = selectionFromConfig(*cfg) + return nil + }) + if err != nil { + return ProviderSelection{}, err + } + + return selection, nil +} + func (s *Service) validate() error { if s == nil || s.manager == nil { return ErrServiceManagerNil @@ -189,7 +229,6 @@ type modelQueryOptions struct { } func (s *Service) modelsForProvider(ctx context.Context, providerCfg config.ProviderConfig, options modelQueryOptions) []ModelDescriptor { - manualModels := modelDescriptorsFromIDs(providerCfg.ConfiguredModels()) defaultModels := modelDescriptorsFromIDs([]string{providerCfg.Model}) cached, cachedOK := s.loadCatalogModels(ctx, providerCfg) @@ -205,7 +244,7 @@ func (s *Service) modelsForProvider(ctx context.Context, providerCfg config.Prov s.queueRefresh(providerCfg) } - return MergeModelDescriptors(manualModels, cached, defaultModels) + return MergeModelDescriptors(cached, defaultModels) } func (s *Service) catalogExpired(ctx context.Context, providerCfg config.ProviderConfig) bool { @@ -314,11 +353,18 @@ func selectionFromConfig(cfg config.Config) ProviderSelection { } } -func selectModel(currentModel string, models []string, fallback string) string { - if config.ContainsModelID(models, currentModel) { - return strings.TrimSpace(currentModel) +func resolveCurrentModel(currentModel string, models []ModelDescriptor, fallback string) (string, bool) { + currentModel = strings.TrimSpace(currentModel) + if containsModelDescriptorID(models, currentModel) { + return currentModel, false + } + + fallback = strings.TrimSpace(fallback) + if fallback != "" && containsModelDescriptorID(models, fallback) { + return fallback, currentModel != fallback } - return strings.TrimSpace(fallback) + + return currentModel, false } func catalogItemFromConfig(cfg config.ProviderConfig, models []ModelDescriptor) ProviderCatalogItem { @@ -329,21 +375,16 @@ func catalogItemFromConfig(cfg config.ProviderConfig, models []ModelDescriptor) } } -func modelDescriptorIDs(models []ModelDescriptor) []string { - if len(models) == 0 { - return nil +func containsModelDescriptorID(models []ModelDescriptor, modelID string) bool { + target := config.NormalizeKey(modelID) + if target == "" { + return false } - ids := make([]string, 0, len(models)) for _, model := range models { - if strings.TrimSpace(model.ID) == "" { - continue + if config.NormalizeKey(model.ID) == target { + return true } - ids = append(ids, model.ID) } - return ids -} - -func containsModelDescriptorID(models []ModelDescriptor, modelID string) bool { - return config.ContainsModelID(modelDescriptorIDs(models), modelID) + return false } diff --git a/internal/provider/service_test.go b/internal/provider/service_test.go index 43e393b9..6a74067a 100644 --- a/internal/provider/service_test.go +++ b/internal/provider/service_test.go @@ -25,25 +25,13 @@ func TestNewService(t *testing.T) { } } -func TestServiceListProvidersFiltersUnsupportedDrivers(t *testing.T) { +func TestServiceListProvidersReturnsBuiltinCatalog(t *testing.T) { t.Parallel() manager := config.NewManager(config.NewLoader(t.TempDir(), testDefaultConfig())) if _, err := manager.Load(context.Background()); err != nil { t.Fatalf("Load() error = %v", err) } - if err := manager.Update(context.Background(), func(cfg *config.Config) error { - cfg.Providers = append(cfg.Providers, config.ProviderConfig{ - Name: "unsupported", - Driver: "custom", - BaseURL: "https://example.com/v1", - Model: "custom-model", - APIKeyEnv: "CUSTOM_API_KEY", - }) - return nil - }); err != nil { - t.Fatalf("append provider: %v", err) - } registry := NewRegistry() if err := registry.Register(testDriverDefinition(config.OpenAIName, nil)); err != nil { @@ -55,10 +43,18 @@ func TestServiceListProvidersFiltersUnsupportedDrivers(t *testing.T) { if err != nil { t.Fatalf("ListProviders() error = %v", err) } - if len(items) != 4 { - t.Fatalf("expected 4 supported builtin providers, got %d", len(items)) + expected := config.DefaultProviders() + if len(items) != len(expected) { + t.Fatalf("expected %d builtin providers, got %d", len(expected), len(items)) + } + expectedNames := make(map[string]struct{}, len(expected)) + for _, providerCfg := range expected { + expectedNames[providerCfg.Name] = struct{}{} } for _, item := range items { + if _, ok := expectedNames[item.ID]; !ok { + t.Fatalf("unexpected provider in catalog: %+v", item) + } if len(item.Models) != 1 { t.Fatalf("expected default fallback model for %q, got %+v", item.ID, item.Models) } @@ -68,17 +64,7 @@ func TestServiceListProvidersFiltersUnsupportedDrivers(t *testing.T) { func TestServiceSelectProviderFallsBackToProviderDefault(t *testing.T) { t.Parallel() - defaults := testDefaultConfig() - defaults.Providers = append(defaults.Providers, config.ProviderConfig{ - Name: "custom-main", - Driver: "custom", - BaseURL: "https://example.com/v1", - Model: "custom-model", - Models: []string{"custom-model", "custom-alt"}, - APIKeyEnv: "CUSTOM_API_KEY", - }) - - manager := config.NewManager(config.NewLoader(t.TempDir(), defaults)) + manager := config.NewManager(config.NewLoader(t.TempDir(), testDefaultConfig())) if _, err := manager.Load(context.Background()); err != nil { t.Fatalf("Load() error = %v", err) } @@ -90,21 +76,21 @@ func TestServiceSelectProviderFallsBackToProviderDefault(t *testing.T) { } registry := NewRegistry() - if err := registry.Register(testDriverDefinition("custom", nil)); err != nil { - t.Fatalf("register custom driver: %v", err) + if err := registry.Register(testDriverDefinition(config.OpenAIName, nil)); err != nil { + t.Fatalf("register openai driver: %v", err) } service := NewService(manager, registry, newMemoryCatalogStore()) - selection, err := service.SelectProvider(context.Background(), "custom-main") + selection, err := service.SelectProvider(context.Background(), config.QiniuName) if err != nil { t.Fatalf("SelectProvider() error = %v", err) } - if selection.ProviderID != "custom-main" || selection.ModelID != "custom-model" { + if selection.ProviderID != config.QiniuName || selection.ModelID != config.QiniuDefaultModel { t.Fatalf("unexpected selection: %+v", selection) } } -func TestServiceListModelsUsesConfiguredModelsWithoutDiscovery(t *testing.T) { +func TestServiceEnsureSelectionRepairsInvalidCurrentModel(t *testing.T) { t.Parallel() manager := config.NewManager(config.NewLoader(t.TempDir(), testDefaultConfig())) @@ -112,19 +98,41 @@ func TestServiceListModelsUsesConfiguredModelsWithoutDiscovery(t *testing.T) { t.Fatalf("Load() error = %v", err) } if err := manager.Update(context.Background(), func(cfg *config.Config) error { - cfg.Providers = append(cfg.Providers, config.ProviderConfig{ - Name: "broken-provider", - Driver: "missing-driver", - BaseURL: "https://example.com/v1", - Model: "broken-model", - Models: []string{"broken-model", "broken-alt"}, - APIKeyEnv: "BROKEN_API_KEY", - }) - cfg.SelectedProvider = "broken-provider" - cfg.CurrentModel = "broken-model" + cfg.CurrentModel = "unsupported-current" return nil }); err != nil { - t.Fatalf("append provider: %v", err) + t.Fatalf("seed invalid current model: %v", err) + } + + registry := NewRegistry() + if err := registry.Register(testDriverDefinition(config.OpenAIName, nil)); err != nil { + t.Fatalf("register openai driver: %v", err) + } + + service := NewService(manager, registry, newMemoryCatalogStore()) + selection, err := service.EnsureSelection(context.Background()) + if err != nil { + t.Fatalf("EnsureSelection() error = %v", err) + } + if selection.ProviderID != config.OpenAIName || selection.ModelID != config.OpenAIDefaultModel { + t.Fatalf("unexpected normalized selection: %+v", selection) + } + + reloaded, err := manager.Reload(context.Background()) + if err != nil { + t.Fatalf("Reload() error = %v", err) + } + if reloaded.CurrentModel != config.OpenAIDefaultModel { + t.Fatalf("expected rewritten current model %q, got %q", config.OpenAIDefaultModel, reloaded.CurrentModel) + } +} + +func TestServiceListModelsFallsBackToDefaultModelWithoutDiscovery(t *testing.T) { + t.Parallel() + + manager := config.NewManager(config.NewLoader(t.TempDir(), testDefaultConfig())) + if _, err := manager.Load(context.Background()); err != nil { + t.Fatalf("Load() error = %v", err) } service := NewService(manager, NewRegistry(), newMemoryCatalogStore()) @@ -132,26 +140,30 @@ func TestServiceListModelsUsesConfiguredModelsWithoutDiscovery(t *testing.T) { if err != nil { t.Fatalf("ListModels() error = %v", err) } - if len(models) != 2 || models[1].ID != "broken-alt" { - t.Fatalf("expected configured models fallback, got %+v", models) + if len(models) != 1 || models[0].ID != config.OpenAIDefaultModel { + t.Fatalf("expected default model fallback, got %+v", models) } } -func TestServiceSetCurrentModelUsesConfiguredModels(t *testing.T) { - t.Parallel() +func TestServiceSetCurrentModelUsesDiscoveredModels(t *testing.T) { + t.Setenv(testAPIKeyEnv, "test-key") manager := config.NewManager(config.NewLoader(t.TempDir(), testDefaultConfig())) if _, err := manager.Load(context.Background()); err != nil { t.Fatalf("Load() error = %v", err) } - if err := manager.Update(context.Background(), func(cfg *config.Config) error { - cfg.Providers[0].Models = []string{cfg.Providers[0].Model, "gpt-4o"} - return nil - }); err != nil { - t.Fatalf("seed provider models: %v", err) + + registry := NewRegistry() + if err := registry.Register(testDriverDefinition(config.OpenAIName, func(ctx context.Context, cfg config.ResolvedProviderConfig) ([]ModelDescriptor, error) { + return []ModelDescriptor{ + {ID: cfg.Model, Name: cfg.Model}, + {ID: "gpt-4o", Name: "GPT-4o"}, + }, nil + })); err != nil { + t.Fatalf("register discovery driver: %v", err) } - service := NewService(manager, NewRegistry(), newMemoryCatalogStore()) + service := NewService(manager, registry, newMemoryCatalogStore()) selection, err := service.SetCurrentModel(context.Background(), "gpt-4o") if err != nil { t.Fatalf("SetCurrentModel() error = %v", err) @@ -177,11 +189,6 @@ func TestServiceListModelsDiscoversAndCachesOnMiss(t *testing.T) { Name: "Server Model", ContextWindow: 32000, MaxOutputTokens: 4096, - Metadata: map[string]any{ - "id": "server-model", - "context_window": float64(32000), - "max_output_tokens": float64(4096), - }, }}, nil })); err != nil { t.Fatalf("register discovery driver: %v", err) @@ -210,25 +217,14 @@ func TestServiceListModelsDiscoversAndCachesOnMiss(t *testing.T) { } func TestServiceListModelsReturnsStaleCacheAndRefreshesInBackground(t *testing.T) { - t.Setenv("CUSTOM_API_KEY", "test-key") - - defaults := testDefaultConfig() - defaults.Providers = []config.ProviderConfig{{ - Name: "custom", - Driver: "custom", - BaseURL: "https://example.com/v1", - Model: "fallback-model", - APIKeyEnv: "CUSTOM_API_KEY", - }} - defaults.SelectedProvider = "custom" - defaults.CurrentModel = "fallback-model" - - manager := config.NewManager(config.NewLoader(t.TempDir(), defaults)) + t.Setenv(testAPIKeyEnv, "test-key") + + manager := config.NewManager(config.NewLoader(t.TempDir(), testDefaultConfig())) if _, err := manager.Load(context.Background()); err != nil { t.Fatalf("Load() error = %v", err) } - identity, err := defaults.Providers[0].Identity() + identity, err := config.OpenAIProvider().Identity() if err != nil { t.Fatalf("Identity() error = %v", err) } @@ -250,7 +246,7 @@ func TestServiceListModelsReturnsStaleCacheAndRefreshesInBackground(t *testing.T refreshed := make(chan struct{}, 1) registry := NewRegistry() - if err := registry.Register(testDriverDefinition("custom", func(ctx context.Context, cfg config.ResolvedProviderConfig) ([]ModelDescriptor, error) { + if err := registry.Register(testDriverDefinition(config.OpenAIName, func(ctx context.Context, cfg config.ResolvedProviderConfig) ([]ModelDescriptor, error) { select { case refreshed <- struct{}{}: default: @@ -335,37 +331,65 @@ func TestServiceBuildAndValidate(t *testing.T) { }) } -func TestSelectModelHelper(t *testing.T) { +func TestResolveCurrentModelHelper(t *testing.T) { t.Parallel() tests := []struct { name string currentModel string - models []string + models []ModelDescriptor fallback string expected string + changed bool }{ { name: "current model in list", currentModel: "gpt-4o", - models: []string{"gpt-4.1", "gpt-4o", "gpt-5.4"}, - fallback: "gpt-4.1", - expected: "gpt-4o", + models: []ModelDescriptor{ + {ID: "gpt-4.1"}, + {ID: "gpt-4o"}, + {ID: "gpt-5.4"}, + }, + fallback: "gpt-4.1", + expected: "gpt-4o", + changed: false, }, { - name: "current model not in list", + name: "current model falls back to provider default", currentModel: "unknown-model", - models: []string{"gpt-4.1", "gpt-4o", "gpt-5.4"}, - fallback: "gpt-4.1", - expected: "gpt-4.1", + models: []ModelDescriptor{ + {ID: "gpt-4.1"}, + {ID: "gpt-4o"}, + {ID: "gpt-5.4"}, + }, + fallback: "gpt-4.1", + expected: "gpt-4.1", + changed: true, + }, + { + name: "missing fallback keeps current model unchanged", + currentModel: "unknown-model", + models: []ModelDescriptor{ + {ID: "gpt-4o"}, + }, + fallback: "gpt-4.1", + expected: "unknown-model", + changed: false, }, } for _, tt := range tests { t.Run(tt.name, func(t *testing.T) { t.Parallel() - if got := selectModel(tt.currentModel, tt.models, tt.fallback); got != tt.expected { - t.Fatalf("selectModel() = %q, want %q", got, tt.expected) + got, changed := resolveCurrentModel(tt.currentModel, tt.models, tt.fallback) + if got != tt.expected || changed != tt.changed { + t.Fatalf( + "resolveCurrentModel() = (%q, %v), want (%q, %v)", + got, + changed, + tt.expected, + tt.changed, + ) } }) } diff --git a/internal/provider/types.go b/internal/provider/types.go index a46d6fe7..ee5b8443 100644 --- a/internal/provider/types.go +++ b/internal/provider/types.go @@ -53,7 +53,6 @@ type ModelDescriptor struct { ContextWindow int `json:"context_window,omitempty"` MaxOutputTokens int `json:"max_output_tokens,omitempty"` Capabilities map[string]bool `json:"capabilities,omitempty"` - Metadata map[string]any `json:"metadata,omitempty"` } type ProviderCatalogItem struct { diff --git a/internal/tui/update_test.go b/internal/tui/update_test.go index 5191f621..685490bb 100644 --- a/internal/tui/update_test.go +++ b/internal/tui/update_test.go @@ -940,7 +940,6 @@ func TestAppUpdateAdditionalTransitions(t *testing.T) { func TestAppUpdateModelPickerEnterAppliesSelection(t *testing.T) { manager := newTestConfigManager(t) if err := manager.Update(context.Background(), func(cfg *config.Config) error { - cfg.Providers[0].Models = []string{cfg.Providers[0].Model, "gpt-4o"} cfg.CurrentModel = "unsupported-current" return nil }); err != nil { @@ -982,18 +981,10 @@ func TestAppUpdateModelPickerEnterAppliesSelection(t *testing.T) { func TestAppUpdateProviderPickerEnterAppliesSelection(t *testing.T) { manager := newTestConfigManager(t) if err := manager.Update(context.Background(), func(cfg *config.Config) error { - cfg.Providers = append(cfg.Providers, config.ProviderConfig{ - Name: "openai-alt", - Driver: "openai", - BaseURL: "https://alt.example.com/v1", - Model: "gpt-4o", - Models: []string{"gpt-4o"}, - APIKeyEnv: config.OpenAIDefaultAPIKeyEnv, - }) cfg.CurrentModel = "unsupported-current" return nil }); err != nil { - t.Fatalf("append provider: %v", err) + t.Fatalf("set unsupported current model: %v", err) } runtime := newStubRuntime() @@ -1013,14 +1004,14 @@ func TestAppUpdateProviderPickerEnterAppliesSelection(t *testing.T) { if !ok { continue } - if candidate.id == "openai-alt" { + if candidate.id == config.QiniuName { selectedIndex = idx selected = candidate.id break } } if selectedIndex < 0 { - t.Fatalf("expected provider picker to include openai-alt") + t.Fatalf("expected provider picker to include %s", config.QiniuName) } app.providerPicker.Select(selectedIndex) @@ -1040,7 +1031,7 @@ func TestAppUpdateProviderPickerEnterAppliesSelection(t *testing.T) { if cfg.SelectedProvider != selected { t.Fatalf("expected selected provider %q, got %q", selected, cfg.SelectedProvider) } - if cfg.CurrentModel != "gpt-4o" { + if cfg.CurrentModel != config.QiniuDefaultModel { t.Fatalf("expected current model to follow provider default, got %q", cfg.CurrentModel) } } @@ -1718,11 +1709,55 @@ func newTestConfigManager(t *testing.T) *config.Manager { func newTestProviderService(t *testing.T, manager *config.Manager) *provider.Service { t.Helper() - registry, err := builtin.NewRegistry() + + registry := provider.NewRegistry() + err := registry.Register(provider.DriverDefinition{ + Name: config.OpenAIName, + Build: func(ctx context.Context, cfg config.ResolvedProviderConfig) (provider.Provider, error) { + return tuiTestProvider{}, nil + }, + }) if err != nil { t.Fatalf("register provider drivers: %v", err) } - return provider.NewService(manager, registry, nil) + return provider.NewService(manager, registry, newTUITestCatalogStore()) +} + +type tuiTestProvider struct{} + +func (tuiTestProvider) Chat(ctx context.Context, req provider.ChatRequest, events chan<- provider.StreamEvent) (provider.ChatResponse, error) { + return provider.ChatResponse{}, nil +} + +type tuiTestCatalogStore struct { + catalogs map[string]provider.ModelCatalog +} + +func newTUITestCatalogStore() *tuiTestCatalogStore { + return &tuiTestCatalogStore{ + catalogs: map[string]provider.ModelCatalog{}, + } +} + +func (s *tuiTestCatalogStore) Load(ctx context.Context, identity config.ProviderIdentity) (provider.ModelCatalog, error) { + if err := ctx.Err(); err != nil { + return provider.ModelCatalog{}, err + } + + catalog, ok := s.catalogs[identity.Key()] + if !ok { + return provider.ModelCatalog{}, provider.ErrModelCatalogNotFound + } + return catalog, nil +} + +func (s *tuiTestCatalogStore) Save(ctx context.Context, catalog provider.ModelCatalog) error { + if err := ctx.Err(); err != nil { + return err + } + + s.catalogs[catalog.Identity.Key()] = catalog + return nil } func collectTeaMessages(cmd tea.Cmd) []tea.Msg { From ea7a43a49c91b57fe2fab73a740694744b20635d Mon Sep 17 00:00:00 2001 From: phantom5099 <1011668688@qq.com> Date: Thu, 2 Apr 2026 21:11:05 +0800 Subject: [PATCH 16/23] =?UTF-8?q?fix:=E4=BF=AE=E5=A4=8D=E6=B5=8B=E8=AF=95?= =?UTF-8?q?=E9=94=99=E8=AF=AF?= MIME-Version: 1.0 Content-Type: text/plain; charset=UTF-8 Content-Transfer-Encoding: 8bit --- internal/provider/service.go | 43 ++++++++++++------ internal/provider/service_test.go | 75 ++++++++++++++++++++++++++++--- internal/tui/app.go | 12 ++++- internal/tui/commands.go | 32 +++++++++++-- internal/tui/provider_service.go | 1 + internal/tui/update.go | 69 +++++++++++++++++++++++----- 6 files changed, 196 insertions(+), 36 deletions(-) diff --git a/internal/provider/service.go b/internal/provider/service.go index 0dbded09..16b846cb 100644 --- a/internal/provider/service.go +++ b/internal/provider/service.go @@ -115,16 +115,23 @@ func (s *Service) ListModels(ctx context.Context) ([]ModelDescriptor, error) { return nil, err } - cfg := s.manager.Get() - selected, err := cfg.SelectedProviderConfig() - if err != nil { + return s.selectedProviderModels(ctx, modelQueryOptions{ + allowSyncRefresh: true, + queueRefresh: true, + }) +} + +func (s *Service) ListModelsSnapshot(ctx context.Context) ([]ModelDescriptor, error) { + if err := s.validate(); err != nil { + return nil, err + } + if err := ctx.Err(); err != nil { return nil, err } - return s.modelsForProvider(ctx, selected, modelQueryOptions{ - allowSyncRefresh: true, - queueRefresh: true, - }), nil + return s.selectedProviderModels(ctx, modelQueryOptions{ + queueRefresh: true, + }) } func (s *Service) SetCurrentModel(ctx context.Context, modelID string) (ProviderSelection, error) { @@ -144,8 +151,7 @@ func (s *Service) SetCurrentModel(ctx context.Context, modelID string) (Provider } models := s.modelsForProvider(ctx, selected, modelQueryOptions{ - allowSyncRefresh: true, - queueRefresh: true, + queueRefresh: true, }) if !containsModelDescriptorID(models, modelID) { return ProviderSelection{}, ErrModelNotFound @@ -189,8 +195,7 @@ func (s *Service) EnsureSelection(ctx context.Context) (ProviderSelection, error } models := s.modelsForProvider(ctx, selected, modelQueryOptions{ - allowSyncRefresh: true, - queueRefresh: true, + queueRefresh: true, }) nextModel, changed := resolveCurrentModel(cfgSnapshot.CurrentModel, models, selected.Model) if !changed { @@ -228,6 +233,16 @@ type modelQueryOptions struct { queueRefresh bool } +func (s *Service) selectedProviderModels(ctx context.Context, options modelQueryOptions) ([]ModelDescriptor, error) { + cfg := s.manager.Get() + selected, err := cfg.SelectedProviderConfig() + if err != nil { + return nil, err + } + + return s.modelsForProvider(ctx, selected, options), nil +} + func (s *Service) modelsForProvider(ctx context.Context, providerCfg config.ProviderConfig, options modelQueryOptions) []ModelDescriptor { defaultModels := modelDescriptorsFromIDs([]string{providerCfg.Model}) @@ -240,8 +255,10 @@ func (s *Service) modelsForProvider(ctx context.Context, providerCfg config.Prov } } - if cachedOK && options.queueRefresh && s.catalogExpired(ctx, providerCfg) { - s.queueRefresh(providerCfg) + if options.queueRefresh { + if !cachedOK || s.catalogExpired(ctx, providerCfg) { + s.queueRefresh(providerCfg) + } } return MergeModelDescriptors(cached, defaultModels) diff --git a/internal/provider/service_test.go b/internal/provider/service_test.go index 6a74067a..f9877da2 100644 --- a/internal/provider/service_test.go +++ b/internal/provider/service_test.go @@ -145,7 +145,7 @@ func TestServiceListModelsFallsBackToDefaultModelWithoutDiscovery(t *testing.T) } } -func TestServiceSetCurrentModelUsesDiscoveredModels(t *testing.T) { +func TestServiceListModelsSnapshotReturnsDefaultAndRefreshesInBackgroundOnMiss(t *testing.T) { t.Setenv(testAPIKeyEnv, "test-key") manager := config.NewManager(config.NewLoader(t.TempDir(), testDefaultConfig())) @@ -153,17 +153,80 @@ func TestServiceSetCurrentModelUsesDiscoveredModels(t *testing.T) { t.Fatalf("Load() error = %v", err) } + store := newMemoryCatalogStore() + refreshed := make(chan struct{}, 1) registry := NewRegistry() if err := registry.Register(testDriverDefinition(config.OpenAIName, func(ctx context.Context, cfg config.ResolvedProviderConfig) ([]ModelDescriptor, error) { - return []ModelDescriptor{ - {ID: cfg.Model, Name: cfg.Model}, - {ID: "gpt-4o", Name: "GPT-4o"}, - }, nil + select { + case refreshed <- struct{}{}: + default: + } + return []ModelDescriptor{{ID: "gpt-4o", Name: "GPT-4o"}}, nil })); err != nil { t.Fatalf("register discovery driver: %v", err) } - service := NewService(manager, registry, newMemoryCatalogStore()) + service := NewService(manager, registry, store) + service.backgroundTimeout = time.Second + + models, err := service.ListModelsSnapshot(context.Background()) + if err != nil { + t.Fatalf("ListModelsSnapshot() error = %v", err) + } + if len(models) != 1 || models[0].ID != config.OpenAIDefaultModel { + t.Fatalf("expected immediate default model fallback, got %+v", models) + } + + select { + case <-refreshed: + case <-time.After(2 * time.Second): + t.Fatal("expected background refresh to run") + } + + identity, err := config.OpenAIProvider().Identity() + if err != nil { + t.Fatalf("Identity() error = %v", err) + } + + deadline := time.Now().Add(2 * time.Second) + for time.Now().Before(deadline) { + catalog, err := store.Load(context.Background(), identity) + if err == nil && containsModelDescriptorID(catalog.Models, "gpt-4o") { + return + } + time.Sleep(20 * time.Millisecond) + } + + catalog, err := store.Load(context.Background(), identity) + if err != nil { + t.Fatalf("Load() refreshed catalog error = %v", err) + } + t.Fatalf("expected refreshed catalog to contain gpt-4o, got %+v", catalog.Models) +} + +func TestServiceSetCurrentModelUsesCachedDiscoveredModels(t *testing.T) { + manager := config.NewManager(config.NewLoader(t.TempDir(), testDefaultConfig())) + if _, err := manager.Load(context.Background()); err != nil { + t.Fatalf("Load() error = %v", err) + } + + store := newMemoryCatalogStore() + identity, err := config.OpenAIProvider().Identity() + if err != nil { + t.Fatalf("Identity() error = %v", err) + } + if err := store.Save(context.Background(), ModelCatalog{ + SchemaVersion: modelCatalogSchemaVersion, + Identity: identity, + Models: []ModelDescriptor{ + {ID: config.OpenAIDefaultModel, Name: config.OpenAIDefaultModel}, + {ID: "gpt-4o", Name: "GPT-4o"}, + }, + }); err != nil { + t.Fatalf("seed model catalog: %v", err) + } + + service := NewService(manager, NewRegistry(), store) selection, err := service.SetCurrentModel(context.Background(), "gpt-4o") if err != nil { t.Fatalf("SetCurrentModel() error = %v", err) diff --git a/internal/tui/app.go b/internal/tui/app.go index 08cd4bff..6907c64f 100644 --- a/internal/tui/app.go +++ b/internal/tui/app.go @@ -32,6 +32,7 @@ type App struct { activeMessages []provider.Message activities []activityEntry fileCandidates []string + modelRefreshID string focus panel width int height int @@ -140,6 +141,7 @@ func New(cfg *config.Config, configManager *config.Manager, runtime agentruntime } app.selectCurrentProvider(cfg.SelectedProvider) app.selectCurrentModel(cfg.CurrentModel) + app.modelRefreshID = cfg.SelectedProvider if err := app.refreshFileCandidates(); err != nil { return App{}, err } @@ -148,5 +150,13 @@ func New(cfg *config.Config, configManager *config.Manager, runtime agentruntime } func (a App) Init() tea.Cmd { - return tea.Batch(ListenForRuntimeEvent(a.runtime.Events()), textarea.Blink, a.spinner.Tick) + cmds := []tea.Cmd{ + ListenForRuntimeEvent(a.runtime.Events()), + textarea.Blink, + a.spinner.Tick, + } + if cmd := runModelCatalogRefresh(a.providerSvc, a.modelRefreshID); cmd != nil { + cmds = append(cmds, cmd) + } + return tea.Batch(cmds...) } diff --git a/internal/tui/commands.go b/internal/tui/commands.go index 0cf8989b..7d03360d 100644 --- a/internal/tui/commands.go +++ b/internal/tui/commands.go @@ -170,7 +170,7 @@ func (a *App) refreshProviderPicker() error { } func (a *App) refreshModelPicker() error { - models, err := a.providerSvc.ListModels(context.Background()) + models, err := a.providerSvc.ListModelsSnapshot(context.Background()) if err != nil { return err } @@ -264,7 +264,8 @@ func runProviderSelection(providerSvc ProviderController, providerName string) t return localCommandResultMsg{err: err} } return localCommandResultMsg{ - notice: fmt.Sprintf("[System] Current provider switched to %s.", selection.ProviderID), + notice: fmt.Sprintf("[System] Current provider switched to %s.", selection.ProviderID), + providerChanged: true, } } } @@ -276,7 +277,8 @@ func runModelSelection(providerSvc ProviderController, modelID string) tea.Cmd { return localCommandResultMsg{err: err} } return localCommandResultMsg{ - notice: fmt.Sprintf("[System] Current model switched to %s.", selection.ModelID), + notice: fmt.Sprintf("[System] Current model switched to %s.", selection.ModelID), + modelChanged: true, } } } @@ -284,7 +286,29 @@ func runModelSelection(providerSvc ProviderController, modelID string) tea.Cmd { func runLocalCommand(configManager *config.Manager, providerSvc ProviderController, snapshot statusSnapshot, raw string) tea.Cmd { return func() tea.Msg { notice, err := executeLocalCommand(context.Background(), configManager, providerSvc, snapshot, raw) - return localCommandResultMsg{notice: notice, err: err} + result := localCommandResultMsg{notice: notice, err: err} + if err == nil { + cfg := configManager.Get() + result.providerChanged = !strings.EqualFold(snapshot.CurrentProvider, cfg.SelectedProvider) + result.modelChanged = !strings.EqualFold(snapshot.CurrentModel, cfg.CurrentModel) + } + return result + } +} + +func runModelCatalogRefresh(providerSvc ProviderController, providerID string) tea.Cmd { + providerID = strings.TrimSpace(providerID) + if providerSvc == nil || providerID == "" { + return nil + } + + return func() tea.Msg { + models, err := providerSvc.ListModels(context.Background()) + return modelCatalogRefreshMsg{ + providerID: providerID, + models: models, + err: err, + } } } diff --git a/internal/tui/provider_service.go b/internal/tui/provider_service.go index c2a5ea1e..256bb0ba 100644 --- a/internal/tui/provider_service.go +++ b/internal/tui/provider_service.go @@ -10,5 +10,6 @@ type ProviderController interface { ListProviders(ctx context.Context) ([]provider.ProviderCatalogItem, error) SelectProvider(ctx context.Context, providerID string) (provider.ProviderSelection, error) ListModels(ctx context.Context) ([]provider.ModelDescriptor, error) + ListModelsSnapshot(ctx context.Context) ([]provider.ModelDescriptor, error) SetCurrentModel(ctx context.Context, modelID string) (provider.ProviderSelection, error) } diff --git a/internal/tui/update.go b/internal/tui/update.go index 6b5b5fc1..e450d3fd 100644 --- a/internal/tui/update.go +++ b/internal/tui/update.go @@ -25,9 +25,16 @@ type runFinishedMsg struct{ err error } type compactFinishedMsg struct { err error } +type modelCatalogRefreshMsg struct { + providerID string + models []provider.ModelDescriptor + err error +} type localCommandResultMsg struct { - notice string - err error + notice string + err error + providerChanged bool + modelChanged bool } const ( @@ -102,6 +109,23 @@ func (a App) Update(msg tea.Msg) (tea.Model, tea.Cmd) { a.rebuildTranscript() a.transcript.GotoBottom() return a, tea.Batch(cmds...) + case modelCatalogRefreshMsg: + if strings.EqualFold(a.modelRefreshID, typed.providerID) { + a.modelRefreshID = "" + } + if !strings.EqualFold(strings.TrimSpace(a.state.CurrentProvider), strings.TrimSpace(typed.providerID)) { + return a, tea.Batch(cmds...) + } + if typed.err != nil { + a.appendActivity("provider", "Failed to refresh models", typed.err.Error(), true) + return a, tea.Batch(cmds...) + } + + a.modelPicker = replacePickerItems(a.modelPicker, newModelPicker(typed.models)) + cfg := a.configManager.Get() + a.syncConfigState(cfg) + a.selectCurrentModel(cfg.CurrentModel) + return a, tea.Batch(cmds...) case localCommandResultMsg: if typed.err != nil { a.state.ExecutionError = typed.err.Error() @@ -112,18 +136,26 @@ func (a App) Update(msg tea.Msg) (tea.Model, tea.Cmd) { a.state.StatusText = typed.notice cfg := a.configManager.Get() a.syncConfigState(cfg) - if err := a.refreshProviderPicker(); err != nil { - a.state.ExecutionError = err.Error() - a.state.StatusText = err.Error() - a.appendActivity("system", "Failed to refresh providers", err.Error(), true) - } - if err := a.refreshModelPicker(); err != nil { - a.state.ExecutionError = err.Error() - a.state.StatusText = err.Error() - a.appendActivity("system", "Failed to refresh models", err.Error(), true) - } else { + if typed.providerChanged { + if err := a.refreshProviderPicker(); err != nil { + a.state.ExecutionError = err.Error() + a.state.StatusText = err.Error() + a.appendActivity("system", "Failed to refresh providers", err.Error(), true) + return a, tea.Batch(cmds...) + } + if err := a.refreshModelPicker(); err != nil { + a.state.ExecutionError = err.Error() + a.state.StatusText = err.Error() + a.appendActivity("system", "Failed to refresh models", err.Error(), true) + return a, tea.Batch(cmds...) + } a.selectCurrentProvider(cfg.SelectedProvider) a.selectCurrentModel(cfg.CurrentModel) + if cmd := a.requestModelCatalogRefresh(cfg.SelectedProvider); cmd != nil { + cmds = append(cmds, cmd) + } + } else if typed.modelChanged { + a.selectCurrentModel(cfg.CurrentModel) } a.appendActivity("command", typed.notice, "", false) } @@ -254,6 +286,9 @@ func (a App) updateInputPanel(msg tea.Msg, typed tea.KeyMsg, cmds []tea.Cmd) (te return a, tea.Batch(cmds...) } a.openModelPicker() + if cmd := a.requestModelCatalogRefresh(a.state.CurrentProvider); cmd != nil { + cmds = append(cmds, cmd) + } return a, tea.Batch(cmds...) } @@ -863,6 +898,16 @@ func (a *App) startDraftSession() { a.rebuildTranscript() } +func (a *App) requestModelCatalogRefresh(providerID string) tea.Cmd { + providerID = strings.TrimSpace(providerID) + if providerID == "" || strings.EqualFold(a.modelRefreshID, providerID) { + return nil + } + + a.modelRefreshID = providerID + return runModelCatalogRefresh(a.providerSvc, providerID) +} + func ListenForRuntimeEvent(sub <-chan agentruntime.RuntimeEvent) tea.Cmd { return func() tea.Msg { event, ok := <-sub From 1ae85ea0623a9d097ae97b963976bd2e642c8df1 Mon Sep 17 00:00:00 2001 From: phantom5099 <1011668688@qq.com> Date: Thu, 2 Apr 2026 22:37:24 +0800 Subject: [PATCH 17/23] =?UTF-8?q?refactor=EF=BC=9A=E4=BC=98=E5=8C=96?= =?UTF-8?q?=E4=B8=80=E4=BA=9B=E5=86=97=E4=BD=99=E4=BB=A3=E7=A0=81?= MIME-Version: 1.0 Content-Type: text/plain; charset=UTF-8 Content-Transfer-Encoding: 8bit --- docs/guides/adding-providers.md | 8 +- internal/config/builtin_providers.go | 44 ++-- internal/config/model.go | 24 +- internal/config/provider_identity.go | 5 + internal/provider/discovery/discovery.go | 2 +- internal/provider/errors.go | 4 +- internal/provider/metadata.go | 6 +- internal/provider/model_catalog_store.go | 10 +- internal/provider/openai/openai.go | 7 - internal/provider/openai/openai_test.go | 9 - internal/provider/provider.go | 1 + internal/provider/registry.go | 78 +++++- internal/provider/registry_test.go | 81 +++++- internal/provider/service.go | 299 +++++++++++++++------- internal/provider/service_test.go | 173 ++++++++++++- internal/provider/transport/retry.go | 48 ++-- internal/provider/transport/retry_test.go | 137 ++++++++-- internal/runtime/runtime_test.go | 4 + internal/tui/commands.go | 8 +- internal/tui/update.go | 12 +- internal/tui/update_test.go | 26 +- internal/tui/view.go | 10 +- 22 files changed, 748 insertions(+), 248 deletions(-) diff --git a/docs/guides/adding-providers.md b/docs/guides/adding-providers.md index db3a0995..91f70517 100644 --- a/docs/guides/adding-providers.md +++ b/docs/guides/adding-providers.md @@ -134,6 +134,11 @@ func (p *Provider) Chat(ctx context.Context, req domain.ChatRequest, events chan // 4. 返回 domain.ChatResponse } +// DiscoverModels 实现模型发现能力 +func (p *Provider) DiscoverModels(ctx context.Context) ([]domain.ModelDescriptor, error) { + return nil, nil +} + // Driver 返回驱动定义 func Driver() domain.DriverDefinition { return domain.DriverDefinition{ @@ -202,7 +207,7 @@ func DefaultProviders() []ProviderConfig { | 类型 | 位置 | 说明 | |------|------|------| -| `Provider` | `internal/provider/types.go` | 核心接口，定义 `Chat` 方法 | +| `Provider` | `internal/provider/types.go` | 核心接口，定义 `Chat` / `DiscoverModels` 方法 | | `DriverDefinition` | `internal/provider/registry.go` | 驱动定义：`Name` + `Build` 构造函数 | | `Registry` | `internal/provider/registry.go` | 驱动注册中心 | @@ -267,4 +272,3 @@ func DefaultProviders() []ProviderConfig { ``` 所有内置 provider 都通过代码集中注册。模型选择器展示的候选模型由默认模型、动态发现结果和本地缓存共同组成。 - diff --git a/internal/config/builtin_providers.go b/internal/config/builtin_providers.go index 57d1a2eb..4cacea33 100644 --- a/internal/config/builtin_providers.go +++ b/internal/config/builtin_providers.go @@ -25,44 +25,48 @@ const ( // OpenAIProvider returns the builtin OpenAI provider definition. func OpenAIProvider() ProviderConfig { return ProviderConfig{ - Name: OpenAIName, - Driver: "openai", - BaseURL: OpenAIDefaultBaseURL, - Model: OpenAIDefaultModel, - APIKeyEnv: OpenAIDefaultAPIKeyEnv, + Name: OpenAIName, + Description: "OpenAI official API", + Driver: "openai", + BaseURL: OpenAIDefaultBaseURL, + Model: OpenAIDefaultModel, + APIKeyEnv: OpenAIDefaultAPIKeyEnv, } } // GeminiProvider returns the builtin Gemini provider definition. func GeminiProvider() ProviderConfig { return ProviderConfig{ - Name: GeminiName, - Driver: "openai", - BaseURL: GeminiDefaultBaseURL, - Model: GeminiDefaultModel, - APIKeyEnv: GeminiDefaultAPIKeyEnv, + Name: GeminiName, + Description: "Google Gemini API (OpenAI-compatible)", + Driver: "openai", + BaseURL: GeminiDefaultBaseURL, + Model: GeminiDefaultModel, + APIKeyEnv: GeminiDefaultAPIKeyEnv, } } // OpenLLProvider returns the builtin OpenLL provider definition. func OpenLLProvider() ProviderConfig { return ProviderConfig{ - Name: OpenLLName, - Driver: "openai", - BaseURL: OpenLLDefaultBaseURL, - Model: OpenLLDefaultModel, - APIKeyEnv: OpenLLDefaultAPIKeyEnv, + Name: OpenLLName, + Description: "OpenLL aggregated API", + Driver: "openai", + BaseURL: OpenLLDefaultBaseURL, + Model: OpenLLDefaultModel, + APIKeyEnv: OpenLLDefaultAPIKeyEnv, } } // QiniuProvider returns the builtin Qiniu provider definition. func QiniuProvider() ProviderConfig { return ProviderConfig{ - Name: QiniuName, - Driver: "openai", - BaseURL: QiniuDefaultBaseURL, - Model: QiniuDefaultModel, - APIKeyEnv: QiniuDefaultAPIKeyEnv, + Name: QiniuName, + Description: "Qiniu AI API", + Driver: "openai", + BaseURL: QiniuDefaultBaseURL, + Model: QiniuDefaultModel, + APIKeyEnv: QiniuDefaultAPIKeyEnv, } } diff --git a/internal/config/model.go b/internal/config/model.go index 4341e49c..0aa42a53 100644 --- a/internal/config/model.go +++ b/internal/config/model.go @@ -48,11 +48,12 @@ type Config struct { } type ProviderConfig struct { - Name string `yaml:"name"` - Driver string `yaml:"driver"` - BaseURL string `yaml:"base_url"` - Model string `yaml:"model"` - APIKeyEnv string `yaml:"api_key_env"` + Name string `yaml:"name"` + Description string `yaml:"description,omitempty"` + Driver string `yaml:"driver"` + BaseURL string `yaml:"base_url"` + Model string `yaml:"model"` + APIKeyEnv string `yaml:"api_key_env"` } type ResolvedProviderConfig struct { @@ -547,16 +548,3 @@ func containsProviderName(providers []ProviderConfig, name string) bool { } return false } - -func ContainsModelID(models []string, model string) bool { - target := NormalizeKey(model) - if target == "" { - return false - } - for _, candidate := range models { - if NormalizeKey(candidate) == target { - return true - } - } - return false -} diff --git a/internal/config/provider_identity.go b/internal/config/provider_identity.go index 93f01207..8f9b5e58 100644 --- a/internal/config/provider_identity.go +++ b/internal/config/provider_identity.go @@ -21,10 +21,15 @@ func NormalizeKey(s string) string { return strings.ToLower(strings.TrimSpace(s)) } +// NormalizeProviderName 将 provider 名称规范化为小写并去除首尾空白。 +// 语义上表示"用户选择的 provider 的显示名称"，用于配置匹配和去重。 func NormalizeProviderName(name string) string { return NormalizeKey(name) } +// NormalizeProviderDriver 将驱动类型标识规范化为小写并去除首尾空白。 +// 语义上表示"协议驱动类型"（如 openai、anthropic），用于注册表匹配。 +// 当前实现与 NormalizeProviderName 相同，保留独立函数以明确语义区分。 func NormalizeProviderDriver(driver string) string { return NormalizeKey(driver) } diff --git a/internal/provider/discovery/discovery.go b/internal/provider/discovery/discovery.go index 4311117c..c4a93e11 100644 --- a/internal/provider/discovery/discovery.go +++ b/internal/provider/discovery/discovery.go @@ -15,7 +15,7 @@ type openAIModelsResponse struct { // FetchOpenAICompatibleModels fetches raw model objects from an OpenAI-compatible /models endpoint. func FetchOpenAICompatibleModels(ctx context.Context, client *http.Client, baseURL string, apiKey string) ([]map[string]any, error) { - endpoint := strings.TrimRight(strings.TrimSpace(baseURL), "/") + "/models" + endpoint := strings.TrimRight(baseURL, "/") + "/models" req, err := http.NewRequestWithContext(ctx, http.MethodGet, endpoint, nil) if err != nil { return nil, fmt.Errorf("provider discovery: build request: %w", err) diff --git a/internal/provider/errors.go b/internal/provider/errors.go index 2280487a..74dc1e3c 100644 --- a/internal/provider/errors.go +++ b/internal/provider/errors.go @@ -12,8 +12,8 @@ var ( ErrModelNotFound = errors.New("model not found") ErrDriverNotFound = errors.New("provider driver not found") ErrDriverAlreadyRegistered = errors.New("provider: driver already registered") - ErrServiceManagerNil = errors.New("provider: config manager is nil") - ErrServiceRegistryNil = errors.New("provider: registry is nil") + errServiceManagerNil = errors.New("provider: config manager is nil") + errServiceRegistryNil = errors.New("provider: registry is nil") ) type ProviderErrorCode string diff --git a/internal/provider/metadata.go b/internal/provider/metadata.go index 319a87f9..f1e0fdcd 100644 --- a/internal/provider/metadata.go +++ b/internal/provider/metadata.go @@ -61,7 +61,7 @@ func MergeModelDescriptors(sources ...[]ModelDescriptor) []ModelDescriptor { for _, source := range sources { for _, candidate := range source { normalized := normalizeModelDescriptor(candidate) - key := normalizeModelID(normalized.ID) + key := config.NormalizeKey(normalized.ID) if key == "" { continue } @@ -204,7 +204,3 @@ func cloneStringBoolMap(source map[string]bool) map[string]bool { } return cloned } - -func normalizeModelID(id string) string { - return config.NormalizeKey(id) -} diff --git a/internal/provider/model_catalog_store.go b/internal/provider/model_catalog_store.go index bca392b9..0d72922b 100644 --- a/internal/provider/model_catalog_store.go +++ b/internal/provider/model_catalog_store.go @@ -87,12 +87,6 @@ func (s *JSONModelCatalogStore) Save(ctx context.Context, catalog ModelCatalog) normalized := normalizeModelCatalog(catalog) - normalizedIdentity, err := config.NewProviderIdentity(normalized.Identity.Driver, normalized.Identity.BaseURL) - if err != nil { - return fmt.Errorf("provider: normalize model catalog key: %w", err) - } - normalized.Identity = normalizedIdentity - path, err := s.catalogPathFromNormalized(normalized.Identity) if err != nil { return err @@ -120,6 +114,10 @@ func normalizeModelCatalog(catalog ModelCatalog) ModelCatalog { if catalog.SchemaVersion == 0 { catalog.SchemaVersion = modelCatalogSchemaVersion } + // 规范化 identity，确保 Save 和 Load 的规范化路径一致。 + if normalized, err := config.NewProviderIdentity(catalog.Identity.Driver, catalog.Identity.BaseURL); err == nil { + catalog.Identity = normalized + } catalog.Models = MergeModelDescriptors(catalog.Models) return catalog } diff --git a/internal/provider/openai/openai.go b/internal/provider/openai/openai.go index c0f891c3..27b55267 100644 --- a/internal/provider/openai/openai.go +++ b/internal/provider/openai/openai.go @@ -52,13 +52,6 @@ func Driver() domain.DriverDefinition { Build: func(ctx context.Context, cfg config.ResolvedProviderConfig) (domain.Provider, error) { return New(cfg, WithTransport(defaultRetryTransport())) }, - Discover: func(ctx context.Context, cfg config.ResolvedProviderConfig) ([]domain.ModelDescriptor, error) { - provider, err := New(cfg, WithTransport(defaultRetryTransport())) - if err != nil { - return nil, err - } - return provider.DiscoverModels(ctx) - }, } } diff --git a/internal/provider/openai/openai_test.go b/internal/provider/openai/openai_test.go index 89c16c0c..d05e16e9 100644 --- a/internal/provider/openai/openai_test.go +++ b/internal/provider/openai/openai_test.go @@ -498,15 +498,6 @@ func (r *readCloser) Close() error { return nil } -func containsString(items []string, target string) bool { - for _, item := range items { - if item == target { - return true - } - } - return false -} - // --- emitToolCallStart / mergeToolCallDeltas 边界补充测试 --- func TestEmitToolCallStartGuards(t *testing.T) { diff --git a/internal/provider/provider.go b/internal/provider/provider.go index 531dba4f..393ce774 100644 --- a/internal/provider/provider.go +++ b/internal/provider/provider.go @@ -4,6 +4,7 @@ import "context" type Provider interface { Chat(ctx context.Context, req ChatRequest, events chan<- StreamEvent) (ChatResponse, error) + DiscoverModels(ctx context.Context) ([]ModelDescriptor, error) } type StreamEventType string diff --git a/internal/provider/registry.go b/internal/provider/registry.go index f5e17324..d4f24075 100644 --- a/internal/provider/registry.go +++ b/internal/provider/registry.go @@ -2,28 +2,34 @@ package provider import ( "context" + "crypto/sha256" + "encoding/hex" "errors" "fmt" "strings" + "sync" "neo-code/internal/config" ) type Builder func(ctx context.Context, cfg config.ResolvedProviderConfig) (Provider, error) -type DiscoveryFunc func(ctx context.Context, cfg config.ResolvedProviderConfig) ([]ModelDescriptor, error) type DriverDefinition struct { - Name string - Build Builder - Discover DiscoveryFunc + Name string + Build Builder } type Registry struct { - drivers map[string]DriverDefinition + mu sync.Mutex + drivers map[string]DriverDefinition + discoveryProviders map[string]Provider } func NewRegistry() *Registry { - return &Registry{drivers: map[string]DriverDefinition{}} + return &Registry{ + drivers: map[string]DriverDefinition{}, + discoveryProviders: map[string]Provider{}, + } } func (r *Registry) Register(driver DriverDefinition) error { @@ -34,7 +40,7 @@ func (r *Registry) Register(driver DriverDefinition) error { r.ensureDrivers() driver.Name = strings.TrimSpace(driver.Name) - driverType := config.NormalizeKey(driver.Name) + driverType := config.NormalizeProviderDriver(driver.Name) if driverType == "" { return errors.New("provider: driver name is empty") } @@ -57,14 +63,11 @@ func (r *Registry) Build(ctx context.Context, cfg config.ResolvedProviderConfig) } func (r *Registry) DiscoverModels(ctx context.Context, cfg config.ResolvedProviderConfig) ([]ModelDescriptor, error) { - driver, err := r.driver(cfg.Driver) + discoveryProvider, err := r.discoveryProvider(ctx, cfg) if err != nil { return nil, err } - if driver.Discover == nil { - return nil, nil - } - return driver.Discover(ctx, cfg) + return discoveryProvider.DiscoverModels(ctx) } func (r *Registry) Supports(driverType string) bool { @@ -76,7 +79,7 @@ func (r *Registry) driver(driverType string) (DriverDefinition, error) { if r == nil { return DriverDefinition{}, ErrDriverNotFound } - driver, ok := r.drivers[config.NormalizeKey(driverType)] + driver, ok := r.drivers[config.NormalizeProviderDriver(driverType)] if !ok { return DriverDefinition{}, fmt.Errorf("%w: %s", ErrDriverNotFound, strings.TrimSpace(driverType)) } @@ -87,4 +90,53 @@ func (r *Registry) ensureDrivers() { if r.drivers == nil { r.drivers = map[string]DriverDefinition{} } + if r.discoveryProviders == nil { + r.discoveryProviders = map[string]Provider{} + } +} + +func (r *Registry) discoveryProvider(ctx context.Context, cfg config.ResolvedProviderConfig) (Provider, error) { + driver, err := r.driver(cfg.Driver) + if err != nil { + return nil, err + } + + cacheKey, err := discoveryCacheKey(cfg) + if err != nil { + return nil, err + } + + r.mu.Lock() + r.ensureDrivers() + if cached, ok := r.discoveryProviders[cacheKey]; ok { + r.mu.Unlock() + return cached, nil + } + r.mu.Unlock() + + built, err := driver.Build(ctx, cfg) + if err != nil { + return nil, err + } + + r.mu.Lock() + r.ensureDrivers() + if cached, ok := r.discoveryProviders[cacheKey]; ok { + r.mu.Unlock() + return cached, nil + } + r.discoveryProviders[cacheKey] = built + r.mu.Unlock() + + return built, nil +} + +func discoveryCacheKey(cfg config.ResolvedProviderConfig) (string, error) { + identity, err := cfg.Identity() + if err != nil { + return "", err + } + + sum := sha256.Sum256([]byte(identity.Key() + "|" + strings.TrimSpace(cfg.APIKey))) + return identity.Key() + "|" + hex.EncodeToString(sum[:]), nil } diff --git a/internal/provider/registry_test.go b/internal/provider/registry_test.go index 7f3cb9df..b377a4d5 100644 --- a/internal/provider/registry_test.go +++ b/internal/provider/registry_test.go @@ -18,6 +18,10 @@ func (stubProvider) Chat(ctx context.Context, req provider.ChatRequest, events c return provider.ChatResponse{}, nil } +func (stubProvider) DiscoverModels(ctx context.Context) ([]provider.ModelDescriptor, error) { + return nil, nil +} + func stubDriver(driverType string) provider.DriverDefinition { return provider.DriverDefinition{ Name: driverType, @@ -99,6 +103,64 @@ func TestRegistryRejectsDuplicateDriverRegistration(t *testing.T) { } } +func TestRegistryDiscoverModelsBuildsProviderOncePerEndpointAndCredential(t *testing.T) { + t.Parallel() + + registry := provider.NewRegistry() + + var mu sync.Mutex + buildCalls := 0 + discoverCalls := 0 + if err := registry.Register(provider.DriverDefinition{ + Name: "custom", + Build: func(ctx context.Context, cfg config.ResolvedProviderConfig) (provider.Provider, error) { + mu.Lock() + buildCalls++ + mu.Unlock() + return discoverStubProvider{ + discover: func(ctx context.Context) ([]provider.ModelDescriptor, error) { + mu.Lock() + discoverCalls++ + mu.Unlock() + return []provider.ModelDescriptor{{ID: "model-a"}}, nil + }, + }, nil + }, + }); err != nil { + t.Fatalf("register custom driver: %v", err) + } + + resolved := config.ResolvedProviderConfig{ + ProviderConfig: config.ProviderConfig{ + Name: "custom-main", + Driver: "custom", + BaseURL: "https://example.com/v1", + Model: "model-a", + APIKeyEnv: "CUSTOM_API_KEY", + }, + APIKey: "test-key", + } + + for range 3 { + models, err := registry.DiscoverModels(context.Background(), resolved) + if err != nil { + t.Fatalf("DiscoverModels() error = %v", err) + } + if len(models) != 1 || models[0].ID != "model-a" { + t.Fatalf("unexpected models: %+v", models) + } + } + + mu.Lock() + defer mu.Unlock() + if buildCalls != 1 { + t.Fatalf("expected discovery provider to be built once, got %d", buildCalls) + } + if discoverCalls != 3 { + t.Fatalf("expected discover method to run each time, got %d", discoverCalls) + } +} + func TestServiceListProvidersFallsBackToProviderDefaultModel(t *testing.T) { t.Parallel() @@ -124,8 +186,8 @@ func TestServiceListProvidersFallsBackToProviderDefaultModel(t *testing.T) { if !ok { t.Fatalf("unexpected builtin provider %q", item.ID) } - if item.Description != "" { - t.Fatalf("expected provider description to stay empty for hidden metadata, got %q", item.Description) + if item.Description == "" { + t.Fatalf("expected provider %q to have a non-empty description", item.ID) } if len(item.Models) != wantModels { t.Fatalf("expected provider models to fall back to the default model, got %+v", item.Models) @@ -315,6 +377,21 @@ type catalogStoreStub struct { catalogs map[string]provider.ModelCatalog } +type discoverStubProvider struct { + discover func(context.Context) ([]provider.ModelDescriptor, error) +} + +func (p discoverStubProvider) Chat(ctx context.Context, req provider.ChatRequest, events chan<- provider.StreamEvent) (provider.ChatResponse, error) { + return provider.ChatResponse{}, nil +} + +func (p discoverStubProvider) DiscoverModels(ctx context.Context) ([]provider.ModelDescriptor, error) { + if p.discover == nil { + return nil, nil + } + return p.discover(ctx) +} + func newCatalogStoreStub() *catalogStoreStub { return &catalogStoreStub{ catalogs: map[string]provider.ModelCatalog{}, diff --git a/internal/provider/service.go b/internal/provider/service.go index 16b846cb..b007663e 100644 --- a/internal/provider/service.go +++ b/internal/provider/service.go @@ -2,6 +2,8 @@ package provider import ( "context" + "errors" + "fmt" "strings" "sync" "time" @@ -20,10 +22,17 @@ type Service struct { catalogs ModelCatalogStore catalogTTL time.Duration backgroundTimeout time.Duration + refreshLease time.Duration now func() time.Time - refreshMu sync.Mutex - inFlightByID map[string]struct{} + refreshMu sync.Mutex + nextRefreshToken uint64 + inFlightByID map[string]refreshTask +} + +type refreshTask struct { + token uint64 + startedAt time.Time } func NewService(manager *config.Manager, registry *Registry, catalogs ModelCatalogStore) *Service { @@ -37,11 +46,16 @@ func NewService(manager *config.Manager, registry *Registry, catalogs ModelCatal catalogs: catalogs, catalogTTL: defaultModelCatalogTTL, backgroundTimeout: defaultBackgroundRefreshTime, + refreshLease: defaultBackgroundRefreshTime * 2, now: time.Now, - inFlightByID: map[string]struct{}{}, + inFlightByID: map[string]refreshTask{}, } } +var errProviderConfigChanged = errors.New("provider: provider config changed during update") + +const maxOptimisticLockRetries = 10 + func (s *Service) ListProviders(ctx context.Context) ([]ProviderCatalogItem, error) { if err := s.validate(); err != nil { return nil, err @@ -72,39 +86,51 @@ func (s *Service) SelectProvider(ctx context.Context, providerName string) (Prov return ProviderSelection{}, err } - cfgSnapshot := s.manager.Get() - providerCfg, err := cfgSnapshot.ProviderByName(providerName) - if err != nil { - return ProviderSelection{}, ErrProviderNotFound - } - if !s.registry.Supports(providerCfg.Driver) { - return ProviderSelection{}, ErrDriverNotFound - } - - models := s.modelsForProvider(ctx, providerCfg, modelQueryOptions{}) + for { + if err := ctx.Err(); err != nil { + return ProviderSelection{}, err + } - var selection ProviderSelection - err = s.manager.Update(ctx, func(cfg *config.Config) error { - selected, err := cfg.ProviderByName(providerName) + cfgSnapshot := s.manager.Get() + providerCfg, err := cfgSnapshot.ProviderByName(providerName) if err != nil { - return ErrProviderNotFound + return ProviderSelection{}, ErrProviderNotFound } - if !s.registry.Supports(selected.Driver) { - return ErrDriverNotFound + if !s.registry.Supports(providerCfg.Driver) { + return ProviderSelection{}, ErrDriverNotFound } - cfg.SelectedProvider = selected.Name - nextModel, _ := resolveCurrentModel(cfg.CurrentModel, models, selected.Model) - cfg.CurrentModel = nextModel - selection = selectionFromConfig(*cfg) - return nil - }) - if err != nil { - return ProviderSelection{}, err - } + models := s.modelsForProvider(ctx, providerCfg, modelQueryOptions{}) + + var selection ProviderSelection + err = s.manager.Update(ctx, func(cfg *config.Config) error { + selected, err := cfg.ProviderByName(providerName) + if err != nil { + return ErrProviderNotFound + } + if selected != providerCfg { + return errProviderConfigChanged + } + if !s.registry.Supports(selected.Driver) { + return ErrDriverNotFound + } + + cfg.SelectedProvider = selected.Name + nextModel, _ := resolveCurrentModel(cfg.CurrentModel, models, selected.Model) + cfg.CurrentModel = nextModel + selection = selectionFromConfig(*cfg) + return nil + }) + if errors.Is(err, errProviderConfigChanged) { + continue + } + if err != nil { + return ProviderSelection{}, err + } - s.queueRefresh(providerCfg) - return selection, nil + s.queueRefresh(providerCfg) + return selection, nil + } } func (s *Service) ListModels(ctx context.Context) ([]ModelDescriptor, error) { @@ -144,33 +170,48 @@ func (s *Service) SetCurrentModel(ctx context.Context, modelID string) (Provider return ProviderSelection{}, ErrModelNotFound } - cfgSnapshot := s.manager.Get() - selected, err := cfgSnapshot.SelectedProviderConfig() - if err != nil { - return ProviderSelection{}, err - } + for attempt := 0; attempt < maxOptimisticLockRetries; attempt++ { + if err := ctx.Err(); err != nil { + return ProviderSelection{}, err + } - models := s.modelsForProvider(ctx, selected, modelQueryOptions{ - queueRefresh: true, - }) - if !containsModelDescriptorID(models, modelID) { - return ProviderSelection{}, ErrModelNotFound - } + cfgSnapshot := s.manager.Get() + selected, err := cfgSnapshot.SelectedProviderConfig() + if err != nil { + return ProviderSelection{}, err + } - var selection ProviderSelection - err = s.manager.Update(ctx, func(cfg *config.Config) error { - if _, err := cfg.SelectedProviderConfig(); err != nil { - return err + models := s.modelsForProvider(ctx, selected, modelQueryOptions{ + queueRefresh: true, + }) + if !containsModelDescriptorID(models, modelID) { + return ProviderSelection{}, ErrModelNotFound } - cfg.CurrentModel = modelID - selection = selectionFromConfig(*cfg) - return nil - }) - if err != nil { - return ProviderSelection{}, err - } - return selection, nil + var selection ProviderSelection + err = s.manager.Update(ctx, func(cfg *config.Config) error { + currentSelected, err := cfg.SelectedProviderConfig() + if err != nil { + return err + } + if currentSelected != selected { + return errProviderConfigChanged + } + + cfg.CurrentModel = modelID + selection = selectionFromConfig(*cfg) + return nil + }) + if errors.Is(err, errProviderConfigChanged) { + continue + } + if err != nil { + return ProviderSelection{}, err + } + + return selection, nil + } + return ProviderSelection{}, fmt.Errorf("provider: set current model %q: config changed too frequently, exhausted %d retries", modelID, maxOptimisticLockRetries) } func (s *Service) Build(ctx context.Context, cfg config.ResolvedProviderConfig) (Provider, error) { @@ -188,42 +229,66 @@ func (s *Service) EnsureSelection(ctx context.Context) (ProviderSelection, error return ProviderSelection{}, err } - cfgSnapshot := s.manager.Get() - selected, err := cfgSnapshot.SelectedProviderConfig() - if err != nil { - return ProviderSelection{}, err - } + for { + cfgSnapshot := s.manager.Get() + selected, err := cfgSnapshot.SelectedProviderConfig() + if err != nil { + return ProviderSelection{}, err + } - models := s.modelsForProvider(ctx, selected, modelQueryOptions{ - queueRefresh: true, - }) - nextModel, changed := resolveCurrentModel(cfgSnapshot.CurrentModel, models, selected.Model) - if !changed { - return selectionFromConfig(cfgSnapshot), nil - } + models := s.modelsForProvider(ctx, selected, modelQueryOptions{ + queueRefresh: true, + }) + nextModel, changed := resolveCurrentModel(cfgSnapshot.CurrentModel, models, selected.Model) + if !changed { + latest := s.manager.Get() + latestSelected, err := latest.SelectedProviderConfig() + if err != nil { + return ProviderSelection{}, err + } + if latestSelected != selected || latest.CurrentModel != cfgSnapshot.CurrentModel { + if err := ctx.Err(); err != nil { + return ProviderSelection{}, err + } + continue + } + return selectionFromConfig(latest), nil + } - var selection ProviderSelection - err = s.manager.Update(ctx, func(cfg *config.Config) error { - if _, err := cfg.SelectedProviderConfig(); err != nil { - return err + var selection ProviderSelection + err = s.manager.Update(ctx, func(cfg *config.Config) error { + currentSelected, err := cfg.SelectedProviderConfig() + if err != nil { + return err + } + if currentSelected != selected { + return errProviderConfigChanged + } + + cfg.CurrentModel = nextModel + selection = selectionFromConfig(*cfg) + return nil + }) + if errors.Is(err, errProviderConfigChanged) { + if err := ctx.Err(); err != nil { + return ProviderSelection{}, err + } + continue + } + if err != nil { + return ProviderSelection{}, err } - cfg.CurrentModel = nextModel - selection = selectionFromConfig(*cfg) - return nil - }) - if err != nil { - return ProviderSelection{}, err - } - return selection, nil + return selection, nil + } } func (s *Service) validate() error { if s == nil || s.manager == nil { - return ErrServiceManagerNil + return errServiceManagerNil } if s.registry == nil { - return ErrServiceRegistryNil + return errServiceRegistryNil } return nil } @@ -342,25 +407,33 @@ func (s *Service) queueRefresh(providerCfg config.ProviderConfig) { } key := identity.Key() + now := s.now() + lease := s.refreshLeaseDuration() + s.refreshMu.Lock() - if _, exists := s.inFlightByID[key]; exists { - s.refreshMu.Unlock() - return + if task, exists := s.inFlightByID[key]; exists { + if lease <= 0 || now.Sub(task.startedAt) < lease { + s.refreshMu.Unlock() + return + } + } + s.nextRefreshToken++ + token := s.nextRefreshToken + s.inFlightByID[key] = refreshTask{ + token: token, + startedAt: now, } - s.inFlightByID[key] = struct{}{} s.refreshMu.Unlock() - go func() { + go func(expectedToken uint64) { defer func() { - s.refreshMu.Lock() - delete(s.inFlightByID, key) - s.refreshMu.Unlock() + s.releaseRefresh(key, expectedToken) }() ctx, cancel := context.WithTimeout(context.Background(), s.backgroundTimeout) defer cancel() _, _ = s.discoverAndPersist(ctx, providerCfg) - }() + }(token) } func selectionFromConfig(cfg config.Config) ProviderSelection { @@ -381,27 +454,65 @@ func resolveCurrentModel(currentModel string, models []ModelDescriptor, fallback return fallback, currentModel != fallback } + if first := firstModelDescriptorID(models); first != "" { + return first, currentModel != first + } + if fallback != "" { + return fallback, currentModel != fallback + } + return currentModel, false } func catalogItemFromConfig(cfg config.ProviderConfig, models []ModelDescriptor) ProviderCatalogItem { return ProviderCatalogItem{ - ID: strings.TrimSpace(cfg.Name), - Name: strings.TrimSpace(cfg.Name), - Models: MergeModelDescriptors(models), + ID: strings.TrimSpace(cfg.Name), + Name: strings.TrimSpace(cfg.Name), + Description: strings.TrimSpace(cfg.Description), + Models: MergeModelDescriptors(models), } } func containsModelDescriptorID(models []ModelDescriptor, modelID string) bool { - target := config.NormalizeKey(modelID) - if target == "" { + if modelID == "" { return false } for _, model := range models { - if config.NormalizeKey(model.ID) == target { + if config.NormalizeKey(model.ID) == modelID { return true } } return false } + +func firstModelDescriptorID(models []ModelDescriptor) string { + for _, model := range models { + id := strings.TrimSpace(model.ID) + if id != "" { + return id + } + } + return "" +} + +func (s *Service) refreshLeaseDuration() time.Duration { + if s.refreshLease > 0 { + return s.refreshLease + } + if s.backgroundTimeout > 0 { + return s.backgroundTimeout * 2 + } + return defaultBackgroundRefreshTime * 2 +} + +func (s *Service) releaseRefresh(key string, expectedToken uint64) { + s.refreshMu.Lock() + defer s.refreshMu.Unlock() + + task, exists := s.inFlightByID[key] + if !exists || task.token != expectedToken { + return + } + delete(s.inFlightByID, key) +} diff --git a/internal/provider/service_test.go b/internal/provider/service_test.go index f9877da2..bfa77067 100644 --- a/internal/provider/service_test.go +++ b/internal/provider/service_test.go @@ -2,6 +2,7 @@ package provider import ( "context" + "errors" "sync" "testing" "time" @@ -90,6 +91,71 @@ func TestServiceSelectProviderFallsBackToProviderDefault(t *testing.T) { } } +func TestServiceSetCurrentModelRetriesWhenSelectedProviderChangesDuringUpdate(t *testing.T) { + t.Parallel() + + manager := config.NewManager(config.NewLoader(t.TempDir(), testDefaultConfig())) + if _, err := manager.Load(context.Background()); err != nil { + t.Fatalf("Load() error = %v", err) + } + + registry := NewRegistry() + if err := registry.Register(testDriverDefinition(config.OpenAIName, nil)); err != nil { + t.Fatalf("register openai driver: %v", err) + } + + store := newBlockingCatalogStore() + openAIIdentity, err := config.OpenAIProvider().Identity() + if err != nil { + t.Fatalf("openai identity: %v", err) + } + if err := store.Save(context.Background(), ModelCatalog{ + SchemaVersion: modelCatalogSchemaVersion, + Identity: openAIIdentity, + Models: []ModelDescriptor{ + {ID: config.OpenAIDefaultModel, Name: config.OpenAIDefaultModel}, + {ID: "gpt-4o", Name: "gpt-4o"}, + }, + }); err != nil { + t.Fatalf("seed openai model catalog: %v", err) + } + + service := NewService(manager, registry, store) + + type result struct { + selection ProviderSelection + err error + } + resultCh := make(chan result, 1) + go func() { + selection, err := service.SetCurrentModel(context.Background(), "gpt-4o") + resultCh <- result{selection: selection, err: err} + }() + + <-store.loadStarted + if err := manager.Update(context.Background(), func(cfg *config.Config) error { + cfg.SelectedProvider = config.QiniuName + cfg.CurrentModel = config.QiniuDefaultModel + return nil + }); err != nil { + t.Fatalf("switch selected provider: %v", err) + } + close(store.releaseLoad) + + outcome := <-resultCh + if !errors.Is(outcome.err, ErrModelNotFound) { + t.Fatalf("expected ErrModelNotFound after provider switch, got selection=%+v err=%v", outcome.selection, outcome.err) + } + + reloaded, err := manager.Reload(context.Background()) + if err != nil { + t.Fatalf("Reload() error = %v", err) + } + if reloaded.SelectedProvider != config.QiniuName || reloaded.CurrentModel != config.QiniuDefaultModel { + t.Fatalf("expected qiniu selection to stay intact, got provider=%q model=%q", reloaded.SelectedProvider, reloaded.CurrentModel) + } +} + func TestServiceEnsureSelectionRepairsInvalidCurrentModel(t *testing.T) { t.Parallel() @@ -353,6 +419,48 @@ func TestServiceListModelsReturnsStaleCacheAndRefreshesInBackground(t *testing.T t.Fatalf("expected refreshed catalog to contain fresh-model, got %+v", catalog.Models) } +func TestServiceQueueRefreshReplacesStaleInFlightTask(t *testing.T) { + t.Setenv(testAPIKeyEnv, "test-key") + + manager := config.NewManager(config.NewLoader(t.TempDir(), testDefaultConfig())) + if _, err := manager.Load(context.Background()); err != nil { + t.Fatalf("Load() error = %v", err) + } + + refreshed := make(chan struct{}, 1) + registry := NewRegistry() + if err := registry.Register(testDriverDefinition(config.OpenAIName, func(ctx context.Context, cfg config.ResolvedProviderConfig) ([]ModelDescriptor, error) { + select { + case refreshed <- struct{}{}: + default: + } + return nil, nil + })); err != nil { + t.Fatalf("register discovery driver: %v", err) + } + + service := NewService(manager, registry, newMemoryCatalogStore()) + service.backgroundTimeout = 100 * time.Millisecond + service.refreshLease = 10 * time.Millisecond + + identity, err := config.OpenAIProvider().Identity() + if err != nil { + t.Fatalf("Identity() error = %v", err) + } + service.inFlightByID[identity.Key()] = refreshTask{ + token: 41, + startedAt: service.now().Add(-time.Minute), + } + + service.queueRefresh(config.OpenAIProvider()) + + select { + case <-refreshed: + case <-time.After(time.Second): + t.Fatal("expected stale in-flight refresh entry to be replaced") + } +} + func TestServiceBuildAndValidate(t *testing.T) { t.Parallel() @@ -430,14 +538,21 @@ func TestResolveCurrentModelHelper(t *testing.T) { changed: true, }, { - name: "missing fallback keeps current model unchanged", + name: "missing fallback uses first discovered model", currentModel: "unknown-model", models: []ModelDescriptor{ {ID: "gpt-4o"}, }, fallback: "gpt-4.1", - expected: "unknown-model", - changed: false, + expected: "gpt-4o", + changed: true, + }, + { + name: "empty models fall back to provider default", + currentModel: "unknown-model", + fallback: "gpt-4.1", + expected: "gpt-4.1", + changed: true, }, } @@ -481,22 +596,37 @@ func testDefaultConfig() *config.Config { return cfg } -func testDriverDefinition(name string, discover DiscoveryFunc) DriverDefinition { +func testDriverDefinition(name string, discover func(context.Context, config.ResolvedProviderConfig) ([]ModelDescriptor, error)) DriverDefinition { return DriverDefinition{ - Name: name, - Discover: discover, + Name: name, Build: func(ctx context.Context, cfg config.ResolvedProviderConfig) (Provider, error) { - return serviceTestProvider{}, nil + return serviceTestProvider{ + discover: func(discoverCtx context.Context) ([]ModelDescriptor, error) { + if discover == nil { + return nil, nil + } + return discover(discoverCtx, cfg) + }, + }, nil }, } } -type serviceTestProvider struct{} +type serviceTestProvider struct { + discover func(context.Context) ([]ModelDescriptor, error) +} func (serviceTestProvider) Chat(ctx context.Context, req ChatRequest, events chan<- StreamEvent) (ChatResponse, error) { return ChatResponse{}, nil } +func (p serviceTestProvider) DiscoverModels(ctx context.Context) ([]ModelDescriptor, error) { + if p.discover == nil { + return nil, nil + } + return p.discover(ctx) +} + type memoryCatalogStore struct { mu sync.Mutex catalogs map[string]ModelCatalog @@ -508,6 +638,33 @@ func newMemoryCatalogStore() *memoryCatalogStore { } } +type blockingCatalogStore struct { + base *memoryCatalogStore + loadStarted chan struct{} + releaseLoad chan struct{} + once sync.Once +} + +func newBlockingCatalogStore() *blockingCatalogStore { + return &blockingCatalogStore{ + base: newMemoryCatalogStore(), + loadStarted: make(chan struct{}), + releaseLoad: make(chan struct{}), + } +} + +func (s *blockingCatalogStore) Load(ctx context.Context, identity config.ProviderIdentity) (ModelCatalog, error) { + s.once.Do(func() { + close(s.loadStarted) + <-s.releaseLoad + }) + return s.base.Load(ctx, identity) +} + +func (s *blockingCatalogStore) Save(ctx context.Context, catalog ModelCatalog) error { + return s.base.Save(ctx, catalog) +} + func (s *memoryCatalogStore) Load(ctx context.Context, identity config.ProviderIdentity) (ModelCatalog, error) { if err := ctx.Err(); err != nil { return ModelCatalog{}, err diff --git a/internal/provider/transport/retry.go b/internal/provider/transport/retry.go index dca6a23b..6b21b83f 100644 --- a/internal/provider/transport/retry.go +++ b/internal/provider/transport/retry.go @@ -5,8 +5,8 @@ import ( "io" "log" "math/rand/v2" + "net" "net/http" - "strings" "time" "neo-code/internal/provider" @@ -57,21 +57,14 @@ func isNetworkError(err error) bool { if err == nil { return false } - msg := err.Error() - // 常见网络错误关键字 - networkIndicators := []string{ - "connection refused", - "connection reset", - "connection timed out", - "timeout", - "no such host", - "i/o timeout", - "dial tcp", + var netErr *net.OpError + if errors.As(err, &netErr) { + return true } - for _, indicator := range networkIndicators { - if strings.Contains(msg, indicator) { - return true - } + // net/url.Error 包装了底层错误，也需要识别。 + var urlErr *net.AddrError + if errors.As(err, &urlErr) { + return true } return false } @@ -126,11 +119,32 @@ func (rt *RetryTransport) RoundTrip(req *http.Request) (*http.Response, error) { resp, err := rt.base.RoundTrip(req) if err == nil { - return resp, nil + if !provider.IsRetryableStatus(resp.StatusCode) { + return resp, nil + } + + lastResp = resp + lastErr = provider.NewProviderErrorFromStatus(resp.StatusCode, resp.Status) + if req.Context().Err() != nil { + _ = rt.drainAndClose(resp) + return nil, req.Context().Err() + } + if attempt == rt.config.MaxRetries { + return resp, nil + } + continue } - lastErr = err + // 将网络层错误包装为结构化 ProviderError。 lastResp = nil + var netErr net.Error + if errors.As(err, &netErr) && netErr.Timeout() { + lastErr = provider.NewTimeoutProviderError(err.Error()) + } else if isNetworkError(err) { + lastErr = provider.NewNetworkProviderError(err.Error()) + } else { + lastErr = err + } // 检查是否可重试。 if !rt.config.RetryableFunc(err) { diff --git a/internal/provider/transport/retry_test.go b/internal/provider/transport/retry_test.go index ac5cc7d7..4142efd2 100644 --- a/internal/provider/transport/retry_test.go +++ b/internal/provider/transport/retry_test.go @@ -6,6 +6,7 @@ import ( "errors" "fmt" "io" + "net" "net/http" "strings" "testing" @@ -19,19 +20,22 @@ import ( func TestIsNetworkError(t *testing.T) { t.Parallel() + dialErr := &net.OpError{Op: "dial", Net: "tcp", Addr: &net.TCPAddr{IP: net.IPv4(127, 0, 0, 1), Port: 8080}, Err: errors.New("connection refused")} + timeoutErr := &net.OpError{Op: "dial", Net: "tcp", Addr: &net.TCPAddr{IP: net.IPv4(127, 0, 0, 1), Port: 443}, Err: errors.New("i/o timeout")} + dnsErr := &net.OpError{Op: "dial", Net: "tcp", Addr: &net.TCPAddr{IP: net.IPv4(0, 0, 0, 0), Port: 0}, Err: &net.DNSError{Err: "no such host"}} + addrErr := &net.AddrError{Err: "no such host", Addr: "api.openai.com"} + tests := []struct { name string err error want bool }{ {"nil", nil, false}, - {"connection refused", errors.New("dial tcp 127.0.0.1:8080: connection refused"), true}, - {"connection reset", errors.New("connection reset by peer"), true}, - {"connection timed out", errors.New("connection timed out"), true}, - {"i/o timeout", errors.New("dial tcp: i/o timeout"), true}, - {"no such host", errors.New("dial tcp: lookup api.openai.com: no such host"), true}, - {"dial tcp", errors.New("dial tcp 127.0.0.1:443: operation was refused"), true}, - {"timeout in message", errors.New("request timeout"), true}, + {"connection refused (OpError)", dialErr, true}, + {"connection refused wrapped", fmt.Errorf("wrapped: %w", dialErr), true}, + {"i/o timeout (OpError)", timeoutErr, true}, + {"DNS failure (OpError)", dnsErr, true}, + {"DNS failure (AddrError)", addrErr, true}, {"tls cert error", errors.New("x509: certificate signed by unknown authority"), false}, {"generic error", errors.New("something went wrong"), false}, {"empty message", errors.New(""), false}, @@ -63,7 +67,7 @@ func TestDefaultRetryableFunc(t *testing.T) { {"non-retryable ProviderError", &provider.ProviderError{StatusCode: 401, Code: provider.ErrorCodeAuthFailed, Retryable: false}, false}, {"wrapped retryable ProviderError", fmt.Errorf("wrapped: %w", &provider.ProviderError{StatusCode: 500, Retryable: true}), true}, {"wrapped non-retryable ProviderError", fmt.Errorf("wrapped: %w", &provider.ProviderError{StatusCode: 400, Retryable: false}), false}, - {"network error", errors.New("dial tcp 127.0.0.1:8080: connection refused"), true}, + {"network error (OpError)", &net.OpError{Op: "dial", Net: "tcp", Err: errors.New("connection refused")}, true}, {"non-network non-provider error", errors.New("bad request"), false}, {"empty error message", errors.New(""), false}, } @@ -179,7 +183,9 @@ func TestRoundTrip_RetryThenSuccess(t *testing.T) { wantResp := newSuccessResponse(http.StatusOK, "ok") m := &mockRoundTripper{ fn: []func() (*http.Response, error){ - func() (*http.Response, error) { return nil, errors.New("dial tcp 127.0.0.1:8080: connection refused") }, + func() (*http.Response, error) { + return nil, &net.OpError{Op: "dial", Net: "tcp", Err: errors.New("connection refused")} + }, func() (*http.Response, error) { return wantResp, nil }, }, } @@ -203,10 +209,41 @@ func TestRoundTrip_RetryThenSuccess(t *testing.T) { } } +func TestRoundTrip_RetryableStatusThenSuccess(t *testing.T) { + t.Parallel() + + m := &mockRoundTripper{ + fn: []func() (*http.Response, error){ + func() (*http.Response, error) { + return newSuccessResponse(http.StatusTooManyRequests, "retry later"), nil + }, + func() (*http.Response, error) { return newSuccessResponse(http.StatusOK, "ok"), nil }, + }, + } + + rt := NewRetryTransport(m, RetryConfig{ + MaxRetries: 2, + WaitBase: 1 * time.Millisecond, + MaxWait: 10 * time.Millisecond, + }) + req, _ := http.NewRequest(http.MethodGet, "http://example.com", nil) + resp, err := rt.RoundTrip(req) + + if err != nil { + t.Fatalf("unexpected error: %v", err) + } + if resp.StatusCode != http.StatusOK { + t.Fatalf("status = %d, want 200", resp.StatusCode) + } + if m.callCount != 2 { + t.Fatalf("callCount = %d, want 2", m.callCount) + } +} + func TestRoundTrip_AllRetriesFail(t *testing.T) { t.Parallel() - networkErr := errors.New("connection refused") + networkErr := &net.OpError{Op: "dial", Net: "tcp", Err: errors.New("connection refused")} m := &mockRoundTripper{ fn: []func() (*http.Response, error){ func() (*http.Response, error) { return nil, networkErr }, @@ -226,14 +263,80 @@ func TestRoundTrip_AllRetriesFail(t *testing.T) { if err == nil { t.Fatal("expected error after all retries exhausted") } - if !errors.Is(err, networkErr) { - t.Fatalf("expected networkErr, got %v", err) + var pErr *provider.ProviderError + if !errors.As(err, &pErr) { + t.Fatalf("expected *ProviderError, got %v", err) + } + if pErr.Code != provider.ErrorCodeNetwork { + t.Fatalf("expected network_error code, got %s", pErr.Code) + } + if m.callCount != 3 { + t.Fatalf("callCount = %d, want 3 (1 initial + 2 retries)", m.callCount) + } +} + +func TestRoundTrip_RetryableStatusExhaustedReturnsFinalResponse(t *testing.T) { + t.Parallel() + + m := &mockRoundTripper{ + fn: []func() (*http.Response, error){ + func() (*http.Response, error) { + return newSuccessResponse(http.StatusInternalServerError, "boom-1"), nil + }, + func() (*http.Response, error) { return newSuccessResponse(http.StatusBadGateway, "boom-2"), nil }, + func() (*http.Response, error) { + return newSuccessResponse(http.StatusServiceUnavailable, "boom-3"), nil + }, + }, + } + + rt := NewRetryTransport(m, RetryConfig{ + MaxRetries: 2, + WaitBase: 1 * time.Millisecond, + MaxWait: 10 * time.Millisecond, + }) + req, _ := http.NewRequest(http.MethodGet, "http://example.com", nil) + resp, err := rt.RoundTrip(req) + + if err != nil { + t.Fatalf("unexpected error: %v", err) + } + if resp.StatusCode != http.StatusServiceUnavailable { + t.Fatalf("status = %d, want 503", resp.StatusCode) } if m.callCount != 3 { t.Fatalf("callCount = %d, want 3 (1 initial + 2 retries)", m.callCount) } } +func TestRoundTrip_ClientErrorStatusIsNotRetried(t *testing.T) { + t.Parallel() + + m := &mockRoundTripper{ + fn: []func() (*http.Response, error){ + func() (*http.Response, error) { return newSuccessResponse(http.StatusBadRequest, "bad request"), nil }, + }, + } + + rt := NewRetryTransport(m, RetryConfig{ + MaxRetries: 2, + WaitBase: 1 * time.Millisecond, + MaxWait: 10 * time.Millisecond, + }) + req, _ := http.NewRequest(http.MethodGet, "http://example.com", nil) + resp, err := rt.RoundTrip(req) + + if err != nil { + t.Fatalf("unexpected error: %v", err) + } + if resp.StatusCode != http.StatusBadRequest { + t.Fatalf("status = %d, want 400", resp.StatusCode) + } + if m.callCount != 1 { + t.Fatalf("callCount = %d, want 1", m.callCount) + } +} + func TestRoundTrip_NonRetryableError(t *testing.T) { t.Parallel() @@ -271,7 +374,9 @@ func TestRoundTrip_ContextCanceled(t *testing.T) { m := &mockRoundTripper{ fn: []func() (*http.Response, error){ - func() (*http.Response, error) { return nil, errors.New("connection refused") }, + func() (*http.Response, error) { + return nil, &net.OpError{Op: "dial", Net: "tcp", Err: errors.New("connection refused")} + }, }, } @@ -302,7 +407,7 @@ func TestRoundTrip_ContextCanceledDuringBackoff(t *testing.T) { fn: []func() (*http.Response, error){ func() (*http.Response, error) { callCount++ - return nil, errors.New("connection refused") + return nil, &net.OpError{Op: "dial", Net: "tcp", Err: errors.New("connection refused")} }, func() (*http.Response, error) { // 在第二次调用之前 context 已取消，这里不应被调用 @@ -337,7 +442,7 @@ func TestRoundTrip_ContextCanceledDuringBackoff(t *testing.T) { func TestRoundTrip_ZeroMaxRetries(t *testing.T) { t.Parallel() - networkErr := errors.New("connection refused") + networkErr := &net.OpError{Op: "dial", Net: "tcp", Err: errors.New("connection refused")} m := &mockRoundTripper{ fn: []func() (*http.Response, error){ func() (*http.Response, error) { return nil, networkErr }, @@ -366,7 +471,7 @@ func TestRoundTrip_SeekerBodyRewound(t *testing.T) { m := &mockRoundTripper{ fn: []func() (*http.Response, error){ func() (*http.Response, error) { - return nil, errors.New("connection refused") + return nil, &net.OpError{Op: "dial", Net: "tcp", Err: errors.New("connection refused")} }, func() (*http.Response, error) { return wantResp, nil diff --git a/internal/runtime/runtime_test.go b/internal/runtime/runtime_test.go index 9ceba70e..0cebc146 100644 --- a/internal/runtime/runtime_test.go +++ b/internal/runtime/runtime_test.go @@ -122,6 +122,10 @@ func (p *scriptedProvider) Chat(ctx context.Context, req provider.ChatRequest, e return p.responses[callIndex], nil } +func (p *scriptedProvider) DiscoverModels(ctx context.Context) ([]provider.ModelDescriptor, error) { + return nil, nil +} + type scriptedProviderFactory struct { provider provider.Provider calls int diff --git a/internal/tui/commands.go b/internal/tui/commands.go index 7d03360d..a318fd7a 100644 --- a/internal/tui/commands.go +++ b/internal/tui/commands.go @@ -76,12 +76,8 @@ const ( messageTagAgent = "[ NEO ]" messageTagTool = "[ TOOL ]" - roleUser = "user" - roleAssistant = "assistant" - roleTool = "tool" - roleEvent = "event" - roleError = "error" - roleSystem = "system" + roleEvent = "event" + roleError = "error" ) type slashCommand struct { diff --git a/internal/tui/update.go b/internal/tui/update.go index e450d3fd..1d86b295 100644 --- a/internal/tui/update.go +++ b/internal/tui/update.go @@ -319,7 +319,7 @@ func (a App) updateInputPanel(msg tea.Msg, typed tea.KeyMsg, cmds []tea.Cmd) (te a.state.ExecutionError = "" a.state.StatusText = statusThinking a.state.CurrentTool = "" - a.activeMessages = append(a.activeMessages, provider.Message{Role: roleUser, Content: input}) + a.activeMessages = append(a.activeMessages, provider.Message{Role: provider.RoleUser, Content: input}) a.rebuildTranscript() cmds = append(cmds, runAgent(a.runtime, a.state.ActiveSessionID, input)) return a, tea.Batch(cmds...) @@ -491,7 +491,7 @@ func (a *App) handleRuntimeEvent(event agentruntime.RuntimeEvent) bool { a.state.CurrentTool = "" if payload, ok := event.Payload.(tools.ToolResult); ok { a.activeMessages = append(a.activeMessages, provider.Message{ - Role: roleTool, + Role: provider.RoleTool, Content: payload.Content, IsError: payload.IsError, }) @@ -523,7 +523,7 @@ func (a *App) handleRuntimeEvent(event agentruntime.RuntimeEvent) bool { a.state.StatusText = statusReady } if payload, ok := event.Payload.(provider.Message); ok && strings.TrimSpace(payload.Content) != "" && !a.lastAssistantMatches(payload.Content) { - a.activeMessages = append(a.activeMessages, provider.Message{Role: roleAssistant, Content: payload.Content}) + a.activeMessages = append(a.activeMessages, provider.Message{Role: provider.RoleAssistant, Content: payload.Content}) transcriptDirty = true } case agentruntime.EventRunCanceled: @@ -599,8 +599,8 @@ func (a *App) appendAssistantChunk(chunk string) { return } - if !a.state.StreamingReply || len(a.activeMessages) == 0 || a.activeMessages[len(a.activeMessages)-1].Role != roleAssistant { - a.activeMessages = append(a.activeMessages, provider.Message{Role: roleAssistant, Content: chunk}) + if !a.state.StreamingReply || len(a.activeMessages) == 0 || a.activeMessages[len(a.activeMessages)-1].Role != provider.RoleAssistant { + a.activeMessages = append(a.activeMessages, provider.Message{Role: provider.RoleAssistant, Content: chunk}) a.state.StreamingReply = true return } @@ -646,7 +646,7 @@ func (a *App) lastAssistantMatches(content string) bool { } last := a.activeMessages[len(a.activeMessages)-1] - return last.Role == roleAssistant && strings.TrimSpace(last.Content) == strings.TrimSpace(content) + return last.Role == provider.RoleAssistant && strings.TrimSpace(last.Content) == strings.TrimSpace(content) } func (a *App) handleViewportKeys(vp *viewport.Model, msg tea.KeyMsg) { diff --git a/internal/tui/update_test.go b/internal/tui/update_test.go index 685490bb..dc6fdd75 100644 --- a/internal/tui/update_test.go +++ b/internal/tui/update_test.go @@ -199,7 +199,7 @@ func TestAppUpdateModelPickerAndRuntimeMessages(t *testing.T) { t.Fatalf("expected assistant draft message") } last := app.activeMessages[len(app.activeMessages)-1] - if last.Role != roleAssistant || last.Content != "hello" { + if last.Role != provider.RoleAssistant || last.Content != "hello" { t.Fatalf("unexpected last assistant draft: %+v", last) } }, @@ -214,7 +214,7 @@ func TestAppUpdateModelPickerAndRuntimeMessages(t *testing.T) { Type: agentruntime.EventAgentDone, SessionID: "session-2", Payload: provider.Message{ - Role: roleAssistant, + Role: provider.RoleAssistant, Content: "final", }, }}, @@ -280,7 +280,7 @@ func TestAppUpdateModelPickerAndRuntimeMessages(t *testing.T) { if app.state.StatusText != statusToolError { t.Fatalf("expected status tool error, got %q", app.state.StatusText) } - if len(app.activeMessages) == 0 || app.activeMessages[len(app.activeMessages)-1].Role != roleTool { + if len(app.activeMessages) == 0 || app.activeMessages[len(app.activeMessages)-1].Role != provider.RoleTool { t.Fatalf("expected tool result to stay in transcript, got %+v", app.activeMessages) } if len(app.activities) == 0 || app.activities[len(app.activities)-1].Title != "Tool error" { @@ -317,8 +317,8 @@ func TestAppHelpersAndRenderingSmoke(t *testing.T) { ID: "session-1", Title: "Existing Session", Messages: []provider.Message{ - {Role: roleUser, Content: "hi"}, - {Role: roleAssistant, Content: "hello"}, + {Role: provider.RoleUser, Content: "hi"}, + {Role: provider.RoleAssistant, Content: "hello"}, }, } runtime.sessions = []agentruntime.SessionSummary{ @@ -342,7 +342,7 @@ func TestAppHelpersAndRenderingSmoke(t *testing.T) { if !app.lastAssistantMatches("hello world") { t.Fatalf("expected assistant draft to match") } - app.appendInlineMessage(roleSystem, "notice") + app.appendInlineMessage(provider.RoleSystem, "notice") app.focusNext() app.focusPrev() @@ -471,7 +471,7 @@ func TestAppHelpersAndRenderingSmoke(t *testing.T) { t.Fatalf("expected error message block") } if app.renderMessageBlock(provider.Message{ - Role: roleAssistant, + Role: provider.RoleAssistant, ToolCalls: []provider.ToolCall{ {Name: "filesystem_edit"}, }, @@ -705,7 +705,7 @@ func TestAppUpdateAdditionalTransitions(t *testing.T) { setup: func(t *testing.T, app *App, runtime *stubRuntime, manager *config.Manager) { app.state.ActiveSessionID = "existing" app.state.ActiveSessionTitle = "Existing" - app.activeMessages = []provider.Message{{Role: roleUser, Content: "hello"}} + app.activeMessages = []provider.Message{{Role: provider.RoleUser, Content: "hello"}} }, msg: tea.KeyMsg{Type: tea.KeyCtrlN}, assert: func(t *testing.T, app App, runtime *stubRuntime, manager *config.Manager, msgs []tea.Msg) { @@ -722,7 +722,7 @@ func TestAppUpdateAdditionalTransitions(t *testing.T) { runtime.loads["s1"] = agentruntime.Session{ ID: "s1", Title: "One", - Messages: []provider.Message{{Role: roleAssistant, Content: "loaded"}}, + Messages: []provider.Message{{Role: provider.RoleAssistant, Content: "loaded"}}, } if err := app.refreshSessions(); err != nil { t.Fatalf("refresh sessions: %v", err) @@ -829,7 +829,7 @@ func TestAppUpdateAdditionalTransitions(t *testing.T) { if !app.state.IsAgentRunning { t.Fatalf("expected agent to start running") } - if len(app.activeMessages) == 0 || app.activeMessages[len(app.activeMessages)-1].Role != roleUser { + if len(app.activeMessages) == 0 || app.activeMessages[len(app.activeMessages)-1].Role != provider.RoleUser { t.Fatalf("expected user message appended") } if len(runtime.runInputs) != 1 || runtime.runInputs[0].Content != "inspect repo\nwith details" { @@ -1122,7 +1122,7 @@ func TestAppHandleRuntimeEventAdditionalBranches(t *testing.T) { if app.state.CurrentTool != "" || app.state.StatusText != statusToolFinished { t.Fatalf("unexpected tool success state: %+v", app.state) } - if len(app.activeMessages) == 0 || app.activeMessages[len(app.activeMessages)-1].Role != roleTool { + if len(app.activeMessages) == 0 || app.activeMessages[len(app.activeMessages)-1].Role != provider.RoleTool { t.Fatalf("expected tool result message in transcript, got %+v", app.activeMessages) } if len(app.activities) == 0 || app.activities[len(app.activities)-1].Title != "Completed tool" { @@ -1729,6 +1729,10 @@ func (tuiTestProvider) Chat(ctx context.Context, req provider.ChatRequest, event return provider.ChatResponse{}, nil } +func (tuiTestProvider) DiscoverModels(ctx context.Context) ([]provider.ModelDescriptor, error) { + return nil, nil +} + type tuiTestCatalogStore struct { catalogs map[string]provider.ModelCatalog } diff --git a/internal/tui/view.go b/internal/tui/view.go index 6b31ae7a..a1f1a9c4 100644 --- a/internal/tui/view.go +++ b/internal/tui/view.go @@ -222,7 +222,7 @@ func (a App) renderMessageBlock(message provider.Message, width int) string { return a.styles.inlineNotice.Width(width).Render(" > " + wrapPlain(message.Content, max(16, width-6))) case roleError: return a.styles.inlineError.Width(width).Render(" ! " + wrapPlain(message.Content, max(16, width-6))) - case roleSystem: + case provider.RoleSystem: return a.styles.inlineSystem.Width(width).Render(" - " + wrapPlain(message.Content, max(16, width-6))) } @@ -233,13 +233,13 @@ func (a App) renderMessageBlock(message provider.Message, width int) string { blockAlign := lipgloss.Left switch message.Role { - case roleUser: + case provider.RoleUser: maxMessageWidth = clamp(int(float64(width)*0.68), 24, width) tag = messageTagUser tagStyle = a.styles.messageUserTag bodyStyle = a.styles.messageUserBody blockAlign = lipgloss.Right - case roleTool: + case provider.RoleTool: tag = messageTagTool tagStyle = a.styles.messageToolTag bodyStyle = a.styles.messageToolBody @@ -258,7 +258,7 @@ func (a App) renderMessageBlock(message provider.Message, width int) string { } contentBlock := a.renderMessageContent(content, maxMessageWidth-2, bodyStyle) - if message.Role == roleUser { + if message.Role == provider.RoleUser { contentBlock = lipgloss.PlaceHorizontal(maxMessageWidth, lipgloss.Right, contentBlock) } @@ -268,7 +268,7 @@ func (a App) renderMessageBlock(message provider.Message, width int) string { contentBlock, ) - if message.Role == roleUser { + if message.Role == provider.RoleUser { return lipgloss.PlaceHorizontal(width, lipgloss.Right, block) } return block From d078249df3dca2f1f80138effdecc4c063f604d0 Mon Sep 17 00:00:00 2001 From: minorcell Date: Thu, 2 Apr 2026 23:35:30 +0800 Subject: [PATCH 18/23] =?UTF-8?q?Revert=20"refactor=EF=BC=9A=E4=BC=98?= =?UTF-8?q?=E5=8C=96=E4=B8=80=E4=BA=9B=E5=86=97=E4=BD=99=E4=BB=A3=E7=A0=81?= =?UTF-8?q?"?= MIME-Version: 1.0 Content-Type: text/plain; charset=UTF-8 Content-Transfer-Encoding: 8bit This reverts commit cba03b1efe79a487e33135108baf301c5da26cc9. --- docs/guides/adding-providers.md | 8 +- internal/config/builtin_providers.go | 44 ++-- internal/config/model.go | 24 +- internal/config/provider_identity.go | 5 - internal/provider/discovery/discovery.go | 2 +- internal/provider/errors.go | 4 +- internal/provider/metadata.go | 6 +- internal/provider/model_catalog_store.go | 10 +- internal/provider/openai/openai.go | 7 + internal/provider/openai/openai_test.go | 9 + internal/provider/provider.go | 1 - internal/provider/registry.go | 78 +----- internal/provider/registry_test.go | 81 +----- internal/provider/service.go | 299 +++++++--------------- internal/provider/service_test.go | 173 +------------ internal/provider/transport/retry.go | 48 ++-- internal/provider/transport/retry_test.go | 137 ++-------- internal/runtime/runtime_test.go | 4 - internal/tui/commands.go | 8 +- internal/tui/update.go | 12 +- internal/tui/update_test.go | 26 +- internal/tui/view.go | 10 +- 22 files changed, 248 insertions(+), 748 deletions(-) diff --git a/docs/guides/adding-providers.md b/docs/guides/adding-providers.md index 91f70517..db3a0995 100644 --- a/docs/guides/adding-providers.md +++ b/docs/guides/adding-providers.md @@ -134,11 +134,6 @@ func (p *Provider) Chat(ctx context.Context, req domain.ChatRequest, events chan // 4. 返回 domain.ChatResponse } -// DiscoverModels 实现模型发现能力 -func (p *Provider) DiscoverModels(ctx context.Context) ([]domain.ModelDescriptor, error) { - return nil, nil -} - // Driver 返回驱动定义 func Driver() domain.DriverDefinition { return domain.DriverDefinition{ @@ -207,7 +202,7 @@ func DefaultProviders() []ProviderConfig { | 类型 | 位置 | 说明 | |------|------|------| -| `Provider` | `internal/provider/types.go` | 核心接口，定义 `Chat` / `DiscoverModels` 方法 | +| `Provider` | `internal/provider/types.go` | 核心接口，定义 `Chat` 方法 | | `DriverDefinition` | `internal/provider/registry.go` | 驱动定义：`Name` + `Build` 构造函数 | | `Registry` | `internal/provider/registry.go` | 驱动注册中心 | @@ -272,3 +267,4 @@ func DefaultProviders() []ProviderConfig { ``` 所有内置 provider 都通过代码集中注册。模型选择器展示的候选模型由默认模型、动态发现结果和本地缓存共同组成。 + diff --git a/internal/config/builtin_providers.go b/internal/config/builtin_providers.go index 4cacea33..57d1a2eb 100644 --- a/internal/config/builtin_providers.go +++ b/internal/config/builtin_providers.go @@ -25,48 +25,44 @@ const ( // OpenAIProvider returns the builtin OpenAI provider definition. func OpenAIProvider() ProviderConfig { return ProviderConfig{ - Name: OpenAIName, - Description: "OpenAI official API", - Driver: "openai", - BaseURL: OpenAIDefaultBaseURL, - Model: OpenAIDefaultModel, - APIKeyEnv: OpenAIDefaultAPIKeyEnv, + Name: OpenAIName, + Driver: "openai", + BaseURL: OpenAIDefaultBaseURL, + Model: OpenAIDefaultModel, + APIKeyEnv: OpenAIDefaultAPIKeyEnv, } } // GeminiProvider returns the builtin Gemini provider definition. func GeminiProvider() ProviderConfig { return ProviderConfig{ - Name: GeminiName, - Description: "Google Gemini API (OpenAI-compatible)", - Driver: "openai", - BaseURL: GeminiDefaultBaseURL, - Model: GeminiDefaultModel, - APIKeyEnv: GeminiDefaultAPIKeyEnv, + Name: GeminiName, + Driver: "openai", + BaseURL: GeminiDefaultBaseURL, + Model: GeminiDefaultModel, + APIKeyEnv: GeminiDefaultAPIKeyEnv, } } // OpenLLProvider returns the builtin OpenLL provider definition. func OpenLLProvider() ProviderConfig { return ProviderConfig{ - Name: OpenLLName, - Description: "OpenLL aggregated API", - Driver: "openai", - BaseURL: OpenLLDefaultBaseURL, - Model: OpenLLDefaultModel, - APIKeyEnv: OpenLLDefaultAPIKeyEnv, + Name: OpenLLName, + Driver: "openai", + BaseURL: OpenLLDefaultBaseURL, + Model: OpenLLDefaultModel, + APIKeyEnv: OpenLLDefaultAPIKeyEnv, } } // QiniuProvider returns the builtin Qiniu provider definition. func QiniuProvider() ProviderConfig { return ProviderConfig{ - Name: QiniuName, - Description: "Qiniu AI API", - Driver: "openai", - BaseURL: QiniuDefaultBaseURL, - Model: QiniuDefaultModel, - APIKeyEnv: QiniuDefaultAPIKeyEnv, + Name: QiniuName, + Driver: "openai", + BaseURL: QiniuDefaultBaseURL, + Model: QiniuDefaultModel, + APIKeyEnv: QiniuDefaultAPIKeyEnv, } } diff --git a/internal/config/model.go b/internal/config/model.go index 0aa42a53..4341e49c 100644 --- a/internal/config/model.go +++ b/internal/config/model.go @@ -48,12 +48,11 @@ type Config struct { } type ProviderConfig struct { - Name string `yaml:"name"` - Description string `yaml:"description,omitempty"` - Driver string `yaml:"driver"` - BaseURL string `yaml:"base_url"` - Model string `yaml:"model"` - APIKeyEnv string `yaml:"api_key_env"` + Name string `yaml:"name"` + Driver string `yaml:"driver"` + BaseURL string `yaml:"base_url"` + Model string `yaml:"model"` + APIKeyEnv string `yaml:"api_key_env"` } type ResolvedProviderConfig struct { @@ -548,3 +547,16 @@ func containsProviderName(providers []ProviderConfig, name string) bool { } return false } + +func ContainsModelID(models []string, model string) bool { + target := NormalizeKey(model) + if target == "" { + return false + } + for _, candidate := range models { + if NormalizeKey(candidate) == target { + return true + } + } + return false +} diff --git a/internal/config/provider_identity.go b/internal/config/provider_identity.go index 8f9b5e58..93f01207 100644 --- a/internal/config/provider_identity.go +++ b/internal/config/provider_identity.go @@ -21,15 +21,10 @@ func NormalizeKey(s string) string { return strings.ToLower(strings.TrimSpace(s)) } -// NormalizeProviderName 将 provider 名称规范化为小写并去除首尾空白。 -// 语义上表示"用户选择的 provider 的显示名称"，用于配置匹配和去重。 func NormalizeProviderName(name string) string { return NormalizeKey(name) } -// NormalizeProviderDriver 将驱动类型标识规范化为小写并去除首尾空白。 -// 语义上表示"协议驱动类型"（如 openai、anthropic），用于注册表匹配。 -// 当前实现与 NormalizeProviderName 相同，保留独立函数以明确语义区分。 func NormalizeProviderDriver(driver string) string { return NormalizeKey(driver) } diff --git a/internal/provider/discovery/discovery.go b/internal/provider/discovery/discovery.go index c4a93e11..4311117c 100644 --- a/internal/provider/discovery/discovery.go +++ b/internal/provider/discovery/discovery.go @@ -15,7 +15,7 @@ type openAIModelsResponse struct { // FetchOpenAICompatibleModels fetches raw model objects from an OpenAI-compatible /models endpoint. func FetchOpenAICompatibleModels(ctx context.Context, client *http.Client, baseURL string, apiKey string) ([]map[string]any, error) { - endpoint := strings.TrimRight(baseURL, "/") + "/models" + endpoint := strings.TrimRight(strings.TrimSpace(baseURL), "/") + "/models" req, err := http.NewRequestWithContext(ctx, http.MethodGet, endpoint, nil) if err != nil { return nil, fmt.Errorf("provider discovery: build request: %w", err) diff --git a/internal/provider/errors.go b/internal/provider/errors.go index 74dc1e3c..2280487a 100644 --- a/internal/provider/errors.go +++ b/internal/provider/errors.go @@ -12,8 +12,8 @@ var ( ErrModelNotFound = errors.New("model not found") ErrDriverNotFound = errors.New("provider driver not found") ErrDriverAlreadyRegistered = errors.New("provider: driver already registered") - errServiceManagerNil = errors.New("provider: config manager is nil") - errServiceRegistryNil = errors.New("provider: registry is nil") + ErrServiceManagerNil = errors.New("provider: config manager is nil") + ErrServiceRegistryNil = errors.New("provider: registry is nil") ) type ProviderErrorCode string diff --git a/internal/provider/metadata.go b/internal/provider/metadata.go index f1e0fdcd..319a87f9 100644 --- a/internal/provider/metadata.go +++ b/internal/provider/metadata.go @@ -61,7 +61,7 @@ func MergeModelDescriptors(sources ...[]ModelDescriptor) []ModelDescriptor { for _, source := range sources { for _, candidate := range source { normalized := normalizeModelDescriptor(candidate) - key := config.NormalizeKey(normalized.ID) + key := normalizeModelID(normalized.ID) if key == "" { continue } @@ -204,3 +204,7 @@ func cloneStringBoolMap(source map[string]bool) map[string]bool { } return cloned } + +func normalizeModelID(id string) string { + return config.NormalizeKey(id) +} diff --git a/internal/provider/model_catalog_store.go b/internal/provider/model_catalog_store.go index 0d72922b..bca392b9 100644 --- a/internal/provider/model_catalog_store.go +++ b/internal/provider/model_catalog_store.go @@ -87,6 +87,12 @@ func (s *JSONModelCatalogStore) Save(ctx context.Context, catalog ModelCatalog) normalized := normalizeModelCatalog(catalog) + normalizedIdentity, err := config.NewProviderIdentity(normalized.Identity.Driver, normalized.Identity.BaseURL) + if err != nil { + return fmt.Errorf("provider: normalize model catalog key: %w", err) + } + normalized.Identity = normalizedIdentity + path, err := s.catalogPathFromNormalized(normalized.Identity) if err != nil { return err @@ -114,10 +120,6 @@ func normalizeModelCatalog(catalog ModelCatalog) ModelCatalog { if catalog.SchemaVersion == 0 { catalog.SchemaVersion = modelCatalogSchemaVersion } - // 规范化 identity，确保 Save 和 Load 的规范化路径一致。 - if normalized, err := config.NewProviderIdentity(catalog.Identity.Driver, catalog.Identity.BaseURL); err == nil { - catalog.Identity = normalized - } catalog.Models = MergeModelDescriptors(catalog.Models) return catalog } diff --git a/internal/provider/openai/openai.go b/internal/provider/openai/openai.go index 27b55267..c0f891c3 100644 --- a/internal/provider/openai/openai.go +++ b/internal/provider/openai/openai.go @@ -52,6 +52,13 @@ func Driver() domain.DriverDefinition { Build: func(ctx context.Context, cfg config.ResolvedProviderConfig) (domain.Provider, error) { return New(cfg, WithTransport(defaultRetryTransport())) }, + Discover: func(ctx context.Context, cfg config.ResolvedProviderConfig) ([]domain.ModelDescriptor, error) { + provider, err := New(cfg, WithTransport(defaultRetryTransport())) + if err != nil { + return nil, err + } + return provider.DiscoverModels(ctx) + }, } } diff --git a/internal/provider/openai/openai_test.go b/internal/provider/openai/openai_test.go index d05e16e9..89c16c0c 100644 --- a/internal/provider/openai/openai_test.go +++ b/internal/provider/openai/openai_test.go @@ -498,6 +498,15 @@ func (r *readCloser) Close() error { return nil } +func containsString(items []string, target string) bool { + for _, item := range items { + if item == target { + return true + } + } + return false +} + // --- emitToolCallStart / mergeToolCallDeltas 边界补充测试 --- func TestEmitToolCallStartGuards(t *testing.T) { diff --git a/internal/provider/provider.go b/internal/provider/provider.go index 393ce774..531dba4f 100644 --- a/internal/provider/provider.go +++ b/internal/provider/provider.go @@ -4,7 +4,6 @@ import "context" type Provider interface { Chat(ctx context.Context, req ChatRequest, events chan<- StreamEvent) (ChatResponse, error) - DiscoverModels(ctx context.Context) ([]ModelDescriptor, error) } type StreamEventType string diff --git a/internal/provider/registry.go b/internal/provider/registry.go index d4f24075..f5e17324 100644 --- a/internal/provider/registry.go +++ b/internal/provider/registry.go @@ -2,34 +2,28 @@ package provider import ( "context" - "crypto/sha256" - "encoding/hex" "errors" "fmt" "strings" - "sync" "neo-code/internal/config" ) type Builder func(ctx context.Context, cfg config.ResolvedProviderConfig) (Provider, error) +type DiscoveryFunc func(ctx context.Context, cfg config.ResolvedProviderConfig) ([]ModelDescriptor, error) type DriverDefinition struct { - Name string - Build Builder + Name string + Build Builder + Discover DiscoveryFunc } type Registry struct { - mu sync.Mutex - drivers map[string]DriverDefinition - discoveryProviders map[string]Provider + drivers map[string]DriverDefinition } func NewRegistry() *Registry { - return &Registry{ - drivers: map[string]DriverDefinition{}, - discoveryProviders: map[string]Provider{}, - } + return &Registry{drivers: map[string]DriverDefinition{}} } func (r *Registry) Register(driver DriverDefinition) error { @@ -40,7 +34,7 @@ func (r *Registry) Register(driver DriverDefinition) error { r.ensureDrivers() driver.Name = strings.TrimSpace(driver.Name) - driverType := config.NormalizeProviderDriver(driver.Name) + driverType := config.NormalizeKey(driver.Name) if driverType == "" { return errors.New("provider: driver name is empty") } @@ -63,11 +57,14 @@ func (r *Registry) Build(ctx context.Context, cfg config.ResolvedProviderConfig) } func (r *Registry) DiscoverModels(ctx context.Context, cfg config.ResolvedProviderConfig) ([]ModelDescriptor, error) { - discoveryProvider, err := r.discoveryProvider(ctx, cfg) + driver, err := r.driver(cfg.Driver) if err != nil { return nil, err } - return discoveryProvider.DiscoverModels(ctx) + if driver.Discover == nil { + return nil, nil + } + return driver.Discover(ctx, cfg) } func (r *Registry) Supports(driverType string) bool { @@ -79,7 +76,7 @@ func (r *Registry) driver(driverType string) (DriverDefinition, error) { if r == nil { return DriverDefinition{}, ErrDriverNotFound } - driver, ok := r.drivers[config.NormalizeProviderDriver(driverType)] + driver, ok := r.drivers[config.NormalizeKey(driverType)] if !ok { return DriverDefinition{}, fmt.Errorf("%w: %s", ErrDriverNotFound, strings.TrimSpace(driverType)) } @@ -90,53 +87,4 @@ func (r *Registry) ensureDrivers() { if r.drivers == nil { r.drivers = map[string]DriverDefinition{} } - if r.discoveryProviders == nil { - r.discoveryProviders = map[string]Provider{} - } -} - -func (r *Registry) discoveryProvider(ctx context.Context, cfg config.ResolvedProviderConfig) (Provider, error) { - driver, err := r.driver(cfg.Driver) - if err != nil { - return nil, err - } - - cacheKey, err := discoveryCacheKey(cfg) - if err != nil { - return nil, err - } - - r.mu.Lock() - r.ensureDrivers() - if cached, ok := r.discoveryProviders[cacheKey]; ok { - r.mu.Unlock() - return cached, nil - } - r.mu.Unlock() - - built, err := driver.Build(ctx, cfg) - if err != nil { - return nil, err - } - - r.mu.Lock() - r.ensureDrivers() - if cached, ok := r.discoveryProviders[cacheKey]; ok { - r.mu.Unlock() - return cached, nil - } - r.discoveryProviders[cacheKey] = built - r.mu.Unlock() - - return built, nil -} - -func discoveryCacheKey(cfg config.ResolvedProviderConfig) (string, error) { - identity, err := cfg.Identity() - if err != nil { - return "", err - } - - sum := sha256.Sum256([]byte(identity.Key() + "|" + strings.TrimSpace(cfg.APIKey))) - return identity.Key() + "|" + hex.EncodeToString(sum[:]), nil } diff --git a/internal/provider/registry_test.go b/internal/provider/registry_test.go index b377a4d5..7f3cb9df 100644 --- a/internal/provider/registry_test.go +++ b/internal/provider/registry_test.go @@ -18,10 +18,6 @@ func (stubProvider) Chat(ctx context.Context, req provider.ChatRequest, events c return provider.ChatResponse{}, nil } -func (stubProvider) DiscoverModels(ctx context.Context) ([]provider.ModelDescriptor, error) { - return nil, nil -} - func stubDriver(driverType string) provider.DriverDefinition { return provider.DriverDefinition{ Name: driverType, @@ -103,64 +99,6 @@ func TestRegistryRejectsDuplicateDriverRegistration(t *testing.T) { } } -func TestRegistryDiscoverModelsBuildsProviderOncePerEndpointAndCredential(t *testing.T) { - t.Parallel() - - registry := provider.NewRegistry() - - var mu sync.Mutex - buildCalls := 0 - discoverCalls := 0 - if err := registry.Register(provider.DriverDefinition{ - Name: "custom", - Build: func(ctx context.Context, cfg config.ResolvedProviderConfig) (provider.Provider, error) { - mu.Lock() - buildCalls++ - mu.Unlock() - return discoverStubProvider{ - discover: func(ctx context.Context) ([]provider.ModelDescriptor, error) { - mu.Lock() - discoverCalls++ - mu.Unlock() - return []provider.ModelDescriptor{{ID: "model-a"}}, nil - }, - }, nil - }, - }); err != nil { - t.Fatalf("register custom driver: %v", err) - } - - resolved := config.ResolvedProviderConfig{ - ProviderConfig: config.ProviderConfig{ - Name: "custom-main", - Driver: "custom", - BaseURL: "https://example.com/v1", - Model: "model-a", - APIKeyEnv: "CUSTOM_API_KEY", - }, - APIKey: "test-key", - } - - for range 3 { - models, err := registry.DiscoverModels(context.Background(), resolved) - if err != nil { - t.Fatalf("DiscoverModels() error = %v", err) - } - if len(models) != 1 || models[0].ID != "model-a" { - t.Fatalf("unexpected models: %+v", models) - } - } - - mu.Lock() - defer mu.Unlock() - if buildCalls != 1 { - t.Fatalf("expected discovery provider to be built once, got %d", buildCalls) - } - if discoverCalls != 3 { - t.Fatalf("expected discover method to run each time, got %d", discoverCalls) - } -} - func TestServiceListProvidersFallsBackToProviderDefaultModel(t *testing.T) { t.Parallel() @@ -186,8 +124,8 @@ func TestServiceListProvidersFallsBackToProviderDefaultModel(t *testing.T) { if !ok { t.Fatalf("unexpected builtin provider %q", item.ID) } - if item.Description == "" { - t.Fatalf("expected provider %q to have a non-empty description", item.ID) + if item.Description != "" { + t.Fatalf("expected provider description to stay empty for hidden metadata, got %q", item.Description) } if len(item.Models) != wantModels { t.Fatalf("expected provider models to fall back to the default model, got %+v", item.Models) @@ -377,21 +315,6 @@ type catalogStoreStub struct { catalogs map[string]provider.ModelCatalog } -type discoverStubProvider struct { - discover func(context.Context) ([]provider.ModelDescriptor, error) -} - -func (p discoverStubProvider) Chat(ctx context.Context, req provider.ChatRequest, events chan<- provider.StreamEvent) (provider.ChatResponse, error) { - return provider.ChatResponse{}, nil -} - -func (p discoverStubProvider) DiscoverModels(ctx context.Context) ([]provider.ModelDescriptor, error) { - if p.discover == nil { - return nil, nil - } - return p.discover(ctx) -} - func newCatalogStoreStub() *catalogStoreStub { return &catalogStoreStub{ catalogs: map[string]provider.ModelCatalog{}, diff --git a/internal/provider/service.go b/internal/provider/service.go index b007663e..16b846cb 100644 --- a/internal/provider/service.go +++ b/internal/provider/service.go @@ -2,8 +2,6 @@ package provider import ( "context" - "errors" - "fmt" "strings" "sync" "time" @@ -22,17 +20,10 @@ type Service struct { catalogs ModelCatalogStore catalogTTL time.Duration backgroundTimeout time.Duration - refreshLease time.Duration now func() time.Time - refreshMu sync.Mutex - nextRefreshToken uint64 - inFlightByID map[string]refreshTask -} - -type refreshTask struct { - token uint64 - startedAt time.Time + refreshMu sync.Mutex + inFlightByID map[string]struct{} } func NewService(manager *config.Manager, registry *Registry, catalogs ModelCatalogStore) *Service { @@ -46,16 +37,11 @@ func NewService(manager *config.Manager, registry *Registry, catalogs ModelCatal catalogs: catalogs, catalogTTL: defaultModelCatalogTTL, backgroundTimeout: defaultBackgroundRefreshTime, - refreshLease: defaultBackgroundRefreshTime * 2, now: time.Now, - inFlightByID: map[string]refreshTask{}, + inFlightByID: map[string]struct{}{}, } } -var errProviderConfigChanged = errors.New("provider: provider config changed during update") - -const maxOptimisticLockRetries = 10 - func (s *Service) ListProviders(ctx context.Context) ([]ProviderCatalogItem, error) { if err := s.validate(); err != nil { return nil, err @@ -86,51 +72,39 @@ func (s *Service) SelectProvider(ctx context.Context, providerName string) (Prov return ProviderSelection{}, err } - for { - if err := ctx.Err(); err != nil { - return ProviderSelection{}, err - } - - cfgSnapshot := s.manager.Get() - providerCfg, err := cfgSnapshot.ProviderByName(providerName) - if err != nil { - return ProviderSelection{}, ErrProviderNotFound - } - if !s.registry.Supports(providerCfg.Driver) { - return ProviderSelection{}, ErrDriverNotFound - } + cfgSnapshot := s.manager.Get() + providerCfg, err := cfgSnapshot.ProviderByName(providerName) + if err != nil { + return ProviderSelection{}, ErrProviderNotFound + } + if !s.registry.Supports(providerCfg.Driver) { + return ProviderSelection{}, ErrDriverNotFound + } - models := s.modelsForProvider(ctx, providerCfg, modelQueryOptions{}) + models := s.modelsForProvider(ctx, providerCfg, modelQueryOptions{}) - var selection ProviderSelection - err = s.manager.Update(ctx, func(cfg *config.Config) error { - selected, err := cfg.ProviderByName(providerName) - if err != nil { - return ErrProviderNotFound - } - if selected != providerCfg { - return errProviderConfigChanged - } - if !s.registry.Supports(selected.Driver) { - return ErrDriverNotFound - } - - cfg.SelectedProvider = selected.Name - nextModel, _ := resolveCurrentModel(cfg.CurrentModel, models, selected.Model) - cfg.CurrentModel = nextModel - selection = selectionFromConfig(*cfg) - return nil - }) - if errors.Is(err, errProviderConfigChanged) { - continue - } + var selection ProviderSelection + err = s.manager.Update(ctx, func(cfg *config.Config) error { + selected, err := cfg.ProviderByName(providerName) if err != nil { - return ProviderSelection{}, err + return ErrProviderNotFound + } + if !s.registry.Supports(selected.Driver) { + return ErrDriverNotFound } - s.queueRefresh(providerCfg) - return selection, nil + cfg.SelectedProvider = selected.Name + nextModel, _ := resolveCurrentModel(cfg.CurrentModel, models, selected.Model) + cfg.CurrentModel = nextModel + selection = selectionFromConfig(*cfg) + return nil + }) + if err != nil { + return ProviderSelection{}, err } + + s.queueRefresh(providerCfg) + return selection, nil } func (s *Service) ListModels(ctx context.Context) ([]ModelDescriptor, error) { @@ -170,48 +144,33 @@ func (s *Service) SetCurrentModel(ctx context.Context, modelID string) (Provider return ProviderSelection{}, ErrModelNotFound } - for attempt := 0; attempt < maxOptimisticLockRetries; attempt++ { - if err := ctx.Err(); err != nil { - return ProviderSelection{}, err - } - - cfgSnapshot := s.manager.Get() - selected, err := cfgSnapshot.SelectedProviderConfig() - if err != nil { - return ProviderSelection{}, err - } + cfgSnapshot := s.manager.Get() + selected, err := cfgSnapshot.SelectedProviderConfig() + if err != nil { + return ProviderSelection{}, err + } - models := s.modelsForProvider(ctx, selected, modelQueryOptions{ - queueRefresh: true, - }) - if !containsModelDescriptorID(models, modelID) { - return ProviderSelection{}, ErrModelNotFound - } + models := s.modelsForProvider(ctx, selected, modelQueryOptions{ + queueRefresh: true, + }) + if !containsModelDescriptorID(models, modelID) { + return ProviderSelection{}, ErrModelNotFound + } - var selection ProviderSelection - err = s.manager.Update(ctx, func(cfg *config.Config) error { - currentSelected, err := cfg.SelectedProviderConfig() - if err != nil { - return err - } - if currentSelected != selected { - return errProviderConfigChanged - } - - cfg.CurrentModel = modelID - selection = selectionFromConfig(*cfg) - return nil - }) - if errors.Is(err, errProviderConfigChanged) { - continue - } - if err != nil { - return ProviderSelection{}, err + var selection ProviderSelection + err = s.manager.Update(ctx, func(cfg *config.Config) error { + if _, err := cfg.SelectedProviderConfig(); err != nil { + return err } - - return selection, nil + cfg.CurrentModel = modelID + selection = selectionFromConfig(*cfg) + return nil + }) + if err != nil { + return ProviderSelection{}, err } - return ProviderSelection{}, fmt.Errorf("provider: set current model %q: config changed too frequently, exhausted %d retries", modelID, maxOptimisticLockRetries) + + return selection, nil } func (s *Service) Build(ctx context.Context, cfg config.ResolvedProviderConfig) (Provider, error) { @@ -229,66 +188,42 @@ func (s *Service) EnsureSelection(ctx context.Context) (ProviderSelection, error return ProviderSelection{}, err } - for { - cfgSnapshot := s.manager.Get() - selected, err := cfgSnapshot.SelectedProviderConfig() - if err != nil { - return ProviderSelection{}, err - } + cfgSnapshot := s.manager.Get() + selected, err := cfgSnapshot.SelectedProviderConfig() + if err != nil { + return ProviderSelection{}, err + } - models := s.modelsForProvider(ctx, selected, modelQueryOptions{ - queueRefresh: true, - }) - nextModel, changed := resolveCurrentModel(cfgSnapshot.CurrentModel, models, selected.Model) - if !changed { - latest := s.manager.Get() - latestSelected, err := latest.SelectedProviderConfig() - if err != nil { - return ProviderSelection{}, err - } - if latestSelected != selected || latest.CurrentModel != cfgSnapshot.CurrentModel { - if err := ctx.Err(); err != nil { - return ProviderSelection{}, err - } - continue - } - return selectionFromConfig(latest), nil - } + models := s.modelsForProvider(ctx, selected, modelQueryOptions{ + queueRefresh: true, + }) + nextModel, changed := resolveCurrentModel(cfgSnapshot.CurrentModel, models, selected.Model) + if !changed { + return selectionFromConfig(cfgSnapshot), nil + } - var selection ProviderSelection - err = s.manager.Update(ctx, func(cfg *config.Config) error { - currentSelected, err := cfg.SelectedProviderConfig() - if err != nil { - return err - } - if currentSelected != selected { - return errProviderConfigChanged - } - - cfg.CurrentModel = nextModel - selection = selectionFromConfig(*cfg) - return nil - }) - if errors.Is(err, errProviderConfigChanged) { - if err := ctx.Err(); err != nil { - return ProviderSelection{}, err - } - continue - } - if err != nil { - return ProviderSelection{}, err + var selection ProviderSelection + err = s.manager.Update(ctx, func(cfg *config.Config) error { + if _, err := cfg.SelectedProviderConfig(); err != nil { + return err } - - return selection, nil + cfg.CurrentModel = nextModel + selection = selectionFromConfig(*cfg) + return nil + }) + if err != nil { + return ProviderSelection{}, err } + + return selection, nil } func (s *Service) validate() error { if s == nil || s.manager == nil { - return errServiceManagerNil + return ErrServiceManagerNil } if s.registry == nil { - return errServiceRegistryNil + return ErrServiceRegistryNil } return nil } @@ -407,33 +342,25 @@ func (s *Service) queueRefresh(providerCfg config.ProviderConfig) { } key := identity.Key() - now := s.now() - lease := s.refreshLeaseDuration() - s.refreshMu.Lock() - if task, exists := s.inFlightByID[key]; exists { - if lease <= 0 || now.Sub(task.startedAt) < lease { - s.refreshMu.Unlock() - return - } - } - s.nextRefreshToken++ - token := s.nextRefreshToken - s.inFlightByID[key] = refreshTask{ - token: token, - startedAt: now, + if _, exists := s.inFlightByID[key]; exists { + s.refreshMu.Unlock() + return } + s.inFlightByID[key] = struct{}{} s.refreshMu.Unlock() - go func(expectedToken uint64) { + go func() { defer func() { - s.releaseRefresh(key, expectedToken) + s.refreshMu.Lock() + delete(s.inFlightByID, key) + s.refreshMu.Unlock() }() ctx, cancel := context.WithTimeout(context.Background(), s.backgroundTimeout) defer cancel() _, _ = s.discoverAndPersist(ctx, providerCfg) - }(token) + }() } func selectionFromConfig(cfg config.Config) ProviderSelection { @@ -454,65 +381,27 @@ func resolveCurrentModel(currentModel string, models []ModelDescriptor, fallback return fallback, currentModel != fallback } - if first := firstModelDescriptorID(models); first != "" { - return first, currentModel != first - } - if fallback != "" { - return fallback, currentModel != fallback - } - return currentModel, false } func catalogItemFromConfig(cfg config.ProviderConfig, models []ModelDescriptor) ProviderCatalogItem { return ProviderCatalogItem{ - ID: strings.TrimSpace(cfg.Name), - Name: strings.TrimSpace(cfg.Name), - Description: strings.TrimSpace(cfg.Description), - Models: MergeModelDescriptors(models), + ID: strings.TrimSpace(cfg.Name), + Name: strings.TrimSpace(cfg.Name), + Models: MergeModelDescriptors(models), } } func containsModelDescriptorID(models []ModelDescriptor, modelID string) bool { - if modelID == "" { + target := config.NormalizeKey(modelID) + if target == "" { return false } for _, model := range models { - if config.NormalizeKey(model.ID) == modelID { + if config.NormalizeKey(model.ID) == target { return true } } return false } - -func firstModelDescriptorID(models []ModelDescriptor) string { - for _, model := range models { - id := strings.TrimSpace(model.ID) - if id != "" { - return id - } - } - return "" -} - -func (s *Service) refreshLeaseDuration() time.Duration { - if s.refreshLease > 0 { - return s.refreshLease - } - if s.backgroundTimeout > 0 { - return s.backgroundTimeout * 2 - } - return defaultBackgroundRefreshTime * 2 -} - -func (s *Service) releaseRefresh(key string, expectedToken uint64) { - s.refreshMu.Lock() - defer s.refreshMu.Unlock() - - task, exists := s.inFlightByID[key] - if !exists || task.token != expectedToken { - return - } - delete(s.inFlightByID, key) -} diff --git a/internal/provider/service_test.go b/internal/provider/service_test.go index bfa77067..f9877da2 100644 --- a/internal/provider/service_test.go +++ b/internal/provider/service_test.go @@ -2,7 +2,6 @@ package provider import ( "context" - "errors" "sync" "testing" "time" @@ -91,71 +90,6 @@ func TestServiceSelectProviderFallsBackToProviderDefault(t *testing.T) { } } -func TestServiceSetCurrentModelRetriesWhenSelectedProviderChangesDuringUpdate(t *testing.T) { - t.Parallel() - - manager := config.NewManager(config.NewLoader(t.TempDir(), testDefaultConfig())) - if _, err := manager.Load(context.Background()); err != nil { - t.Fatalf("Load() error = %v", err) - } - - registry := NewRegistry() - if err := registry.Register(testDriverDefinition(config.OpenAIName, nil)); err != nil { - t.Fatalf("register openai driver: %v", err) - } - - store := newBlockingCatalogStore() - openAIIdentity, err := config.OpenAIProvider().Identity() - if err != nil { - t.Fatalf("openai identity: %v", err) - } - if err := store.Save(context.Background(), ModelCatalog{ - SchemaVersion: modelCatalogSchemaVersion, - Identity: openAIIdentity, - Models: []ModelDescriptor{ - {ID: config.OpenAIDefaultModel, Name: config.OpenAIDefaultModel}, - {ID: "gpt-4o", Name: "gpt-4o"}, - }, - }); err != nil { - t.Fatalf("seed openai model catalog: %v", err) - } - - service := NewService(manager, registry, store) - - type result struct { - selection ProviderSelection - err error - } - resultCh := make(chan result, 1) - go func() { - selection, err := service.SetCurrentModel(context.Background(), "gpt-4o") - resultCh <- result{selection: selection, err: err} - }() - - <-store.loadStarted - if err := manager.Update(context.Background(), func(cfg *config.Config) error { - cfg.SelectedProvider = config.QiniuName - cfg.CurrentModel = config.QiniuDefaultModel - return nil - }); err != nil { - t.Fatalf("switch selected provider: %v", err) - } - close(store.releaseLoad) - - outcome := <-resultCh - if !errors.Is(outcome.err, ErrModelNotFound) { - t.Fatalf("expected ErrModelNotFound after provider switch, got selection=%+v err=%v", outcome.selection, outcome.err) - } - - reloaded, err := manager.Reload(context.Background()) - if err != nil { - t.Fatalf("Reload() error = %v", err) - } - if reloaded.SelectedProvider != config.QiniuName || reloaded.CurrentModel != config.QiniuDefaultModel { - t.Fatalf("expected qiniu selection to stay intact, got provider=%q model=%q", reloaded.SelectedProvider, reloaded.CurrentModel) - } -} - func TestServiceEnsureSelectionRepairsInvalidCurrentModel(t *testing.T) { t.Parallel() @@ -419,48 +353,6 @@ func TestServiceListModelsReturnsStaleCacheAndRefreshesInBackground(t *testing.T t.Fatalf("expected refreshed catalog to contain fresh-model, got %+v", catalog.Models) } -func TestServiceQueueRefreshReplacesStaleInFlightTask(t *testing.T) { - t.Setenv(testAPIKeyEnv, "test-key") - - manager := config.NewManager(config.NewLoader(t.TempDir(), testDefaultConfig())) - if _, err := manager.Load(context.Background()); err != nil { - t.Fatalf("Load() error = %v", err) - } - - refreshed := make(chan struct{}, 1) - registry := NewRegistry() - if err := registry.Register(testDriverDefinition(config.OpenAIName, func(ctx context.Context, cfg config.ResolvedProviderConfig) ([]ModelDescriptor, error) { - select { - case refreshed <- struct{}{}: - default: - } - return nil, nil - })); err != nil { - t.Fatalf("register discovery driver: %v", err) - } - - service := NewService(manager, registry, newMemoryCatalogStore()) - service.backgroundTimeout = 100 * time.Millisecond - service.refreshLease = 10 * time.Millisecond - - identity, err := config.OpenAIProvider().Identity() - if err != nil { - t.Fatalf("Identity() error = %v", err) - } - service.inFlightByID[identity.Key()] = refreshTask{ - token: 41, - startedAt: service.now().Add(-time.Minute), - } - - service.queueRefresh(config.OpenAIProvider()) - - select { - case <-refreshed: - case <-time.After(time.Second): - t.Fatal("expected stale in-flight refresh entry to be replaced") - } -} - func TestServiceBuildAndValidate(t *testing.T) { t.Parallel() @@ -538,21 +430,14 @@ func TestResolveCurrentModelHelper(t *testing.T) { changed: true, }, { - name: "missing fallback uses first discovered model", + name: "missing fallback keeps current model unchanged", currentModel: "unknown-model", models: []ModelDescriptor{ {ID: "gpt-4o"}, }, fallback: "gpt-4.1", - expected: "gpt-4o", - changed: true, - }, - { - name: "empty models fall back to provider default", - currentModel: "unknown-model", - fallback: "gpt-4.1", - expected: "gpt-4.1", - changed: true, + expected: "unknown-model", + changed: false, }, } @@ -596,37 +481,22 @@ func testDefaultConfig() *config.Config { return cfg } -func testDriverDefinition(name string, discover func(context.Context, config.ResolvedProviderConfig) ([]ModelDescriptor, error)) DriverDefinition { +func testDriverDefinition(name string, discover DiscoveryFunc) DriverDefinition { return DriverDefinition{ - Name: name, + Name: name, + Discover: discover, Build: func(ctx context.Context, cfg config.ResolvedProviderConfig) (Provider, error) { - return serviceTestProvider{ - discover: func(discoverCtx context.Context) ([]ModelDescriptor, error) { - if discover == nil { - return nil, nil - } - return discover(discoverCtx, cfg) - }, - }, nil + return serviceTestProvider{}, nil }, } } -type serviceTestProvider struct { - discover func(context.Context) ([]ModelDescriptor, error) -} +type serviceTestProvider struct{} func (serviceTestProvider) Chat(ctx context.Context, req ChatRequest, events chan<- StreamEvent) (ChatResponse, error) { return ChatResponse{}, nil } -func (p serviceTestProvider) DiscoverModels(ctx context.Context) ([]ModelDescriptor, error) { - if p.discover == nil { - return nil, nil - } - return p.discover(ctx) -} - type memoryCatalogStore struct { mu sync.Mutex catalogs map[string]ModelCatalog @@ -638,33 +508,6 @@ func newMemoryCatalogStore() *memoryCatalogStore { } } -type blockingCatalogStore struct { - base *memoryCatalogStore - loadStarted chan struct{} - releaseLoad chan struct{} - once sync.Once -} - -func newBlockingCatalogStore() *blockingCatalogStore { - return &blockingCatalogStore{ - base: newMemoryCatalogStore(), - loadStarted: make(chan struct{}), - releaseLoad: make(chan struct{}), - } -} - -func (s *blockingCatalogStore) Load(ctx context.Context, identity config.ProviderIdentity) (ModelCatalog, error) { - s.once.Do(func() { - close(s.loadStarted) - <-s.releaseLoad - }) - return s.base.Load(ctx, identity) -} - -func (s *blockingCatalogStore) Save(ctx context.Context, catalog ModelCatalog) error { - return s.base.Save(ctx, catalog) -} - func (s *memoryCatalogStore) Load(ctx context.Context, identity config.ProviderIdentity) (ModelCatalog, error) { if err := ctx.Err(); err != nil { return ModelCatalog{}, err diff --git a/internal/provider/transport/retry.go b/internal/provider/transport/retry.go index 6b21b83f..dca6a23b 100644 --- a/internal/provider/transport/retry.go +++ b/internal/provider/transport/retry.go @@ -5,8 +5,8 @@ import ( "io" "log" "math/rand/v2" - "net" "net/http" + "strings" "time" "neo-code/internal/provider" @@ -57,14 +57,21 @@ func isNetworkError(err error) bool { if err == nil { return false } - var netErr *net.OpError - if errors.As(err, &netErr) { - return true + msg := err.Error() + // 常见网络错误关键字 + networkIndicators := []string{ + "connection refused", + "connection reset", + "connection timed out", + "timeout", + "no such host", + "i/o timeout", + "dial tcp", } - // net/url.Error 包装了底层错误，也需要识别。 - var urlErr *net.AddrError - if errors.As(err, &urlErr) { - return true + for _, indicator := range networkIndicators { + if strings.Contains(msg, indicator) { + return true + } } return false } @@ -119,32 +126,11 @@ func (rt *RetryTransport) RoundTrip(req *http.Request) (*http.Response, error) { resp, err := rt.base.RoundTrip(req) if err == nil { - if !provider.IsRetryableStatus(resp.StatusCode) { - return resp, nil - } - - lastResp = resp - lastErr = provider.NewProviderErrorFromStatus(resp.StatusCode, resp.Status) - if req.Context().Err() != nil { - _ = rt.drainAndClose(resp) - return nil, req.Context().Err() - } - if attempt == rt.config.MaxRetries { - return resp, nil - } - continue + return resp, nil } - // 将网络层错误包装为结构化 ProviderError。 + lastErr = err lastResp = nil - var netErr net.Error - if errors.As(err, &netErr) && netErr.Timeout() { - lastErr = provider.NewTimeoutProviderError(err.Error()) - } else if isNetworkError(err) { - lastErr = provider.NewNetworkProviderError(err.Error()) - } else { - lastErr = err - } // 检查是否可重试。 if !rt.config.RetryableFunc(err) { diff --git a/internal/provider/transport/retry_test.go b/internal/provider/transport/retry_test.go index 4142efd2..ac5cc7d7 100644 --- a/internal/provider/transport/retry_test.go +++ b/internal/provider/transport/retry_test.go @@ -6,7 +6,6 @@ import ( "errors" "fmt" "io" - "net" "net/http" "strings" "testing" @@ -20,22 +19,19 @@ import ( func TestIsNetworkError(t *testing.T) { t.Parallel() - dialErr := &net.OpError{Op: "dial", Net: "tcp", Addr: &net.TCPAddr{IP: net.IPv4(127, 0, 0, 1), Port: 8080}, Err: errors.New("connection refused")} - timeoutErr := &net.OpError{Op: "dial", Net: "tcp", Addr: &net.TCPAddr{IP: net.IPv4(127, 0, 0, 1), Port: 443}, Err: errors.New("i/o timeout")} - dnsErr := &net.OpError{Op: "dial", Net: "tcp", Addr: &net.TCPAddr{IP: net.IPv4(0, 0, 0, 0), Port: 0}, Err: &net.DNSError{Err: "no such host"}} - addrErr := &net.AddrError{Err: "no such host", Addr: "api.openai.com"} - tests := []struct { name string err error want bool }{ {"nil", nil, false}, - {"connection refused (OpError)", dialErr, true}, - {"connection refused wrapped", fmt.Errorf("wrapped: %w", dialErr), true}, - {"i/o timeout (OpError)", timeoutErr, true}, - {"DNS failure (OpError)", dnsErr, true}, - {"DNS failure (AddrError)", addrErr, true}, + {"connection refused", errors.New("dial tcp 127.0.0.1:8080: connection refused"), true}, + {"connection reset", errors.New("connection reset by peer"), true}, + {"connection timed out", errors.New("connection timed out"), true}, + {"i/o timeout", errors.New("dial tcp: i/o timeout"), true}, + {"no such host", errors.New("dial tcp: lookup api.openai.com: no such host"), true}, + {"dial tcp", errors.New("dial tcp 127.0.0.1:443: operation was refused"), true}, + {"timeout in message", errors.New("request timeout"), true}, {"tls cert error", errors.New("x509: certificate signed by unknown authority"), false}, {"generic error", errors.New("something went wrong"), false}, {"empty message", errors.New(""), false}, @@ -67,7 +63,7 @@ func TestDefaultRetryableFunc(t *testing.T) { {"non-retryable ProviderError", &provider.ProviderError{StatusCode: 401, Code: provider.ErrorCodeAuthFailed, Retryable: false}, false}, {"wrapped retryable ProviderError", fmt.Errorf("wrapped: %w", &provider.ProviderError{StatusCode: 500, Retryable: true}), true}, {"wrapped non-retryable ProviderError", fmt.Errorf("wrapped: %w", &provider.ProviderError{StatusCode: 400, Retryable: false}), false}, - {"network error (OpError)", &net.OpError{Op: "dial", Net: "tcp", Err: errors.New("connection refused")}, true}, + {"network error", errors.New("dial tcp 127.0.0.1:8080: connection refused"), true}, {"non-network non-provider error", errors.New("bad request"), false}, {"empty error message", errors.New(""), false}, } @@ -183,9 +179,7 @@ func TestRoundTrip_RetryThenSuccess(t *testing.T) { wantResp := newSuccessResponse(http.StatusOK, "ok") m := &mockRoundTripper{ fn: []func() (*http.Response, error){ - func() (*http.Response, error) { - return nil, &net.OpError{Op: "dial", Net: "tcp", Err: errors.New("connection refused")} - }, + func() (*http.Response, error) { return nil, errors.New("dial tcp 127.0.0.1:8080: connection refused") }, func() (*http.Response, error) { return wantResp, nil }, }, } @@ -209,41 +203,10 @@ func TestRoundTrip_RetryThenSuccess(t *testing.T) { } } -func TestRoundTrip_RetryableStatusThenSuccess(t *testing.T) { - t.Parallel() - - m := &mockRoundTripper{ - fn: []func() (*http.Response, error){ - func() (*http.Response, error) { - return newSuccessResponse(http.StatusTooManyRequests, "retry later"), nil - }, - func() (*http.Response, error) { return newSuccessResponse(http.StatusOK, "ok"), nil }, - }, - } - - rt := NewRetryTransport(m, RetryConfig{ - MaxRetries: 2, - WaitBase: 1 * time.Millisecond, - MaxWait: 10 * time.Millisecond, - }) - req, _ := http.NewRequest(http.MethodGet, "http://example.com", nil) - resp, err := rt.RoundTrip(req) - - if err != nil { - t.Fatalf("unexpected error: %v", err) - } - if resp.StatusCode != http.StatusOK { - t.Fatalf("status = %d, want 200", resp.StatusCode) - } - if m.callCount != 2 { - t.Fatalf("callCount = %d, want 2", m.callCount) - } -} - func TestRoundTrip_AllRetriesFail(t *testing.T) { t.Parallel() - networkErr := &net.OpError{Op: "dial", Net: "tcp", Err: errors.New("connection refused")} + networkErr := errors.New("connection refused") m := &mockRoundTripper{ fn: []func() (*http.Response, error){ func() (*http.Response, error) { return nil, networkErr }, @@ -263,80 +226,14 @@ func TestRoundTrip_AllRetriesFail(t *testing.T) { if err == nil { t.Fatal("expected error after all retries exhausted") } - var pErr *provider.ProviderError - if !errors.As(err, &pErr) { - t.Fatalf("expected *ProviderError, got %v", err) - } - if pErr.Code != provider.ErrorCodeNetwork { - t.Fatalf("expected network_error code, got %s", pErr.Code) - } - if m.callCount != 3 { - t.Fatalf("callCount = %d, want 3 (1 initial + 2 retries)", m.callCount) - } -} - -func TestRoundTrip_RetryableStatusExhaustedReturnsFinalResponse(t *testing.T) { - t.Parallel() - - m := &mockRoundTripper{ - fn: []func() (*http.Response, error){ - func() (*http.Response, error) { - return newSuccessResponse(http.StatusInternalServerError, "boom-1"), nil - }, - func() (*http.Response, error) { return newSuccessResponse(http.StatusBadGateway, "boom-2"), nil }, - func() (*http.Response, error) { - return newSuccessResponse(http.StatusServiceUnavailable, "boom-3"), nil - }, - }, - } - - rt := NewRetryTransport(m, RetryConfig{ - MaxRetries: 2, - WaitBase: 1 * time.Millisecond, - MaxWait: 10 * time.Millisecond, - }) - req, _ := http.NewRequest(http.MethodGet, "http://example.com", nil) - resp, err := rt.RoundTrip(req) - - if err != nil { - t.Fatalf("unexpected error: %v", err) - } - if resp.StatusCode != http.StatusServiceUnavailable { - t.Fatalf("status = %d, want 503", resp.StatusCode) + if !errors.Is(err, networkErr) { + t.Fatalf("expected networkErr, got %v", err) } if m.callCount != 3 { t.Fatalf("callCount = %d, want 3 (1 initial + 2 retries)", m.callCount) } } -func TestRoundTrip_ClientErrorStatusIsNotRetried(t *testing.T) { - t.Parallel() - - m := &mockRoundTripper{ - fn: []func() (*http.Response, error){ - func() (*http.Response, error) { return newSuccessResponse(http.StatusBadRequest, "bad request"), nil }, - }, - } - - rt := NewRetryTransport(m, RetryConfig{ - MaxRetries: 2, - WaitBase: 1 * time.Millisecond, - MaxWait: 10 * time.Millisecond, - }) - req, _ := http.NewRequest(http.MethodGet, "http://example.com", nil) - resp, err := rt.RoundTrip(req) - - if err != nil { - t.Fatalf("unexpected error: %v", err) - } - if resp.StatusCode != http.StatusBadRequest { - t.Fatalf("status = %d, want 400", resp.StatusCode) - } - if m.callCount != 1 { - t.Fatalf("callCount = %d, want 1", m.callCount) - } -} - func TestRoundTrip_NonRetryableError(t *testing.T) { t.Parallel() @@ -374,9 +271,7 @@ func TestRoundTrip_ContextCanceled(t *testing.T) { m := &mockRoundTripper{ fn: []func() (*http.Response, error){ - func() (*http.Response, error) { - return nil, &net.OpError{Op: "dial", Net: "tcp", Err: errors.New("connection refused")} - }, + func() (*http.Response, error) { return nil, errors.New("connection refused") }, }, } @@ -407,7 +302,7 @@ func TestRoundTrip_ContextCanceledDuringBackoff(t *testing.T) { fn: []func() (*http.Response, error){ func() (*http.Response, error) { callCount++ - return nil, &net.OpError{Op: "dial", Net: "tcp", Err: errors.New("connection refused")} + return nil, errors.New("connection refused") }, func() (*http.Response, error) { // 在第二次调用之前 context 已取消，这里不应被调用 @@ -442,7 +337,7 @@ func TestRoundTrip_ContextCanceledDuringBackoff(t *testing.T) { func TestRoundTrip_ZeroMaxRetries(t *testing.T) { t.Parallel() - networkErr := &net.OpError{Op: "dial", Net: "tcp", Err: errors.New("connection refused")} + networkErr := errors.New("connection refused") m := &mockRoundTripper{ fn: []func() (*http.Response, error){ func() (*http.Response, error) { return nil, networkErr }, @@ -471,7 +366,7 @@ func TestRoundTrip_SeekerBodyRewound(t *testing.T) { m := &mockRoundTripper{ fn: []func() (*http.Response, error){ func() (*http.Response, error) { - return nil, &net.OpError{Op: "dial", Net: "tcp", Err: errors.New("connection refused")} + return nil, errors.New("connection refused") }, func() (*http.Response, error) { return wantResp, nil diff --git a/internal/runtime/runtime_test.go b/internal/runtime/runtime_test.go index 0cebc146..9ceba70e 100644 --- a/internal/runtime/runtime_test.go +++ b/internal/runtime/runtime_test.go @@ -122,10 +122,6 @@ func (p *scriptedProvider) Chat(ctx context.Context, req provider.ChatRequest, e return p.responses[callIndex], nil } -func (p *scriptedProvider) DiscoverModels(ctx context.Context) ([]provider.ModelDescriptor, error) { - return nil, nil -} - type scriptedProviderFactory struct { provider provider.Provider calls int diff --git a/internal/tui/commands.go b/internal/tui/commands.go index a318fd7a..7d03360d 100644 --- a/internal/tui/commands.go +++ b/internal/tui/commands.go @@ -76,8 +76,12 @@ const ( messageTagAgent = "[ NEO ]" messageTagTool = "[ TOOL ]" - roleEvent = "event" - roleError = "error" + roleUser = "user" + roleAssistant = "assistant" + roleTool = "tool" + roleEvent = "event" + roleError = "error" + roleSystem = "system" ) type slashCommand struct { diff --git a/internal/tui/update.go b/internal/tui/update.go index 1d86b295..e450d3fd 100644 --- a/internal/tui/update.go +++ b/internal/tui/update.go @@ -319,7 +319,7 @@ func (a App) updateInputPanel(msg tea.Msg, typed tea.KeyMsg, cmds []tea.Cmd) (te a.state.ExecutionError = "" a.state.StatusText = statusThinking a.state.CurrentTool = "" - a.activeMessages = append(a.activeMessages, provider.Message{Role: provider.RoleUser, Content: input}) + a.activeMessages = append(a.activeMessages, provider.Message{Role: roleUser, Content: input}) a.rebuildTranscript() cmds = append(cmds, runAgent(a.runtime, a.state.ActiveSessionID, input)) return a, tea.Batch(cmds...) @@ -491,7 +491,7 @@ func (a *App) handleRuntimeEvent(event agentruntime.RuntimeEvent) bool { a.state.CurrentTool = "" if payload, ok := event.Payload.(tools.ToolResult); ok { a.activeMessages = append(a.activeMessages, provider.Message{ - Role: provider.RoleTool, + Role: roleTool, Content: payload.Content, IsError: payload.IsError, }) @@ -523,7 +523,7 @@ func (a *App) handleRuntimeEvent(event agentruntime.RuntimeEvent) bool { a.state.StatusText = statusReady } if payload, ok := event.Payload.(provider.Message); ok && strings.TrimSpace(payload.Content) != "" && !a.lastAssistantMatches(payload.Content) { - a.activeMessages = append(a.activeMessages, provider.Message{Role: provider.RoleAssistant, Content: payload.Content}) + a.activeMessages = append(a.activeMessages, provider.Message{Role: roleAssistant, Content: payload.Content}) transcriptDirty = true } case agentruntime.EventRunCanceled: @@ -599,8 +599,8 @@ func (a *App) appendAssistantChunk(chunk string) { return } - if !a.state.StreamingReply || len(a.activeMessages) == 0 || a.activeMessages[len(a.activeMessages)-1].Role != provider.RoleAssistant { - a.activeMessages = append(a.activeMessages, provider.Message{Role: provider.RoleAssistant, Content: chunk}) + if !a.state.StreamingReply || len(a.activeMessages) == 0 || a.activeMessages[len(a.activeMessages)-1].Role != roleAssistant { + a.activeMessages = append(a.activeMessages, provider.Message{Role: roleAssistant, Content: chunk}) a.state.StreamingReply = true return } @@ -646,7 +646,7 @@ func (a *App) lastAssistantMatches(content string) bool { } last := a.activeMessages[len(a.activeMessages)-1] - return last.Role == provider.RoleAssistant && strings.TrimSpace(last.Content) == strings.TrimSpace(content) + return last.Role == roleAssistant && strings.TrimSpace(last.Content) == strings.TrimSpace(content) } func (a *App) handleViewportKeys(vp *viewport.Model, msg tea.KeyMsg) { diff --git a/internal/tui/update_test.go b/internal/tui/update_test.go index dc6fdd75..685490bb 100644 --- a/internal/tui/update_test.go +++ b/internal/tui/update_test.go @@ -199,7 +199,7 @@ func TestAppUpdateModelPickerAndRuntimeMessages(t *testing.T) { t.Fatalf("expected assistant draft message") } last := app.activeMessages[len(app.activeMessages)-1] - if last.Role != provider.RoleAssistant || last.Content != "hello" { + if last.Role != roleAssistant || last.Content != "hello" { t.Fatalf("unexpected last assistant draft: %+v", last) } }, @@ -214,7 +214,7 @@ func TestAppUpdateModelPickerAndRuntimeMessages(t *testing.T) { Type: agentruntime.EventAgentDone, SessionID: "session-2", Payload: provider.Message{ - Role: provider.RoleAssistant, + Role: roleAssistant, Content: "final", }, }}, @@ -280,7 +280,7 @@ func TestAppUpdateModelPickerAndRuntimeMessages(t *testing.T) { if app.state.StatusText != statusToolError { t.Fatalf("expected status tool error, got %q", app.state.StatusText) } - if len(app.activeMessages) == 0 || app.activeMessages[len(app.activeMessages)-1].Role != provider.RoleTool { + if len(app.activeMessages) == 0 || app.activeMessages[len(app.activeMessages)-1].Role != roleTool { t.Fatalf("expected tool result to stay in transcript, got %+v", app.activeMessages) } if len(app.activities) == 0 || app.activities[len(app.activities)-1].Title != "Tool error" { @@ -317,8 +317,8 @@ func TestAppHelpersAndRenderingSmoke(t *testing.T) { ID: "session-1", Title: "Existing Session", Messages: []provider.Message{ - {Role: provider.RoleUser, Content: "hi"}, - {Role: provider.RoleAssistant, Content: "hello"}, + {Role: roleUser, Content: "hi"}, + {Role: roleAssistant, Content: "hello"}, }, } runtime.sessions = []agentruntime.SessionSummary{ @@ -342,7 +342,7 @@ func TestAppHelpersAndRenderingSmoke(t *testing.T) { if !app.lastAssistantMatches("hello world") { t.Fatalf("expected assistant draft to match") } - app.appendInlineMessage(provider.RoleSystem, "notice") + app.appendInlineMessage(roleSystem, "notice") app.focusNext() app.focusPrev() @@ -471,7 +471,7 @@ func TestAppHelpersAndRenderingSmoke(t *testing.T) { t.Fatalf("expected error message block") } if app.renderMessageBlock(provider.Message{ - Role: provider.RoleAssistant, + Role: roleAssistant, ToolCalls: []provider.ToolCall{ {Name: "filesystem_edit"}, }, @@ -705,7 +705,7 @@ func TestAppUpdateAdditionalTransitions(t *testing.T) { setup: func(t *testing.T, app *App, runtime *stubRuntime, manager *config.Manager) { app.state.ActiveSessionID = "existing" app.state.ActiveSessionTitle = "Existing" - app.activeMessages = []provider.Message{{Role: provider.RoleUser, Content: "hello"}} + app.activeMessages = []provider.Message{{Role: roleUser, Content: "hello"}} }, msg: tea.KeyMsg{Type: tea.KeyCtrlN}, assert: func(t *testing.T, app App, runtime *stubRuntime, manager *config.Manager, msgs []tea.Msg) { @@ -722,7 +722,7 @@ func TestAppUpdateAdditionalTransitions(t *testing.T) { runtime.loads["s1"] = agentruntime.Session{ ID: "s1", Title: "One", - Messages: []provider.Message{{Role: provider.RoleAssistant, Content: "loaded"}}, + Messages: []provider.Message{{Role: roleAssistant, Content: "loaded"}}, } if err := app.refreshSessions(); err != nil { t.Fatalf("refresh sessions: %v", err) @@ -829,7 +829,7 @@ func TestAppUpdateAdditionalTransitions(t *testing.T) { if !app.state.IsAgentRunning { t.Fatalf("expected agent to start running") } - if len(app.activeMessages) == 0 || app.activeMessages[len(app.activeMessages)-1].Role != provider.RoleUser { + if len(app.activeMessages) == 0 || app.activeMessages[len(app.activeMessages)-1].Role != roleUser { t.Fatalf("expected user message appended") } if len(runtime.runInputs) != 1 || runtime.runInputs[0].Content != "inspect repo\nwith details" { @@ -1122,7 +1122,7 @@ func TestAppHandleRuntimeEventAdditionalBranches(t *testing.T) { if app.state.CurrentTool != "" || app.state.StatusText != statusToolFinished { t.Fatalf("unexpected tool success state: %+v", app.state) } - if len(app.activeMessages) == 0 || app.activeMessages[len(app.activeMessages)-1].Role != provider.RoleTool { + if len(app.activeMessages) == 0 || app.activeMessages[len(app.activeMessages)-1].Role != roleTool { t.Fatalf("expected tool result message in transcript, got %+v", app.activeMessages) } if len(app.activities) == 0 || app.activities[len(app.activities)-1].Title != "Completed tool" { @@ -1729,10 +1729,6 @@ func (tuiTestProvider) Chat(ctx context.Context, req provider.ChatRequest, event return provider.ChatResponse{}, nil } -func (tuiTestProvider) DiscoverModels(ctx context.Context) ([]provider.ModelDescriptor, error) { - return nil, nil -} - type tuiTestCatalogStore struct { catalogs map[string]provider.ModelCatalog } diff --git a/internal/tui/view.go b/internal/tui/view.go index a1f1a9c4..6b31ae7a 100644 --- a/internal/tui/view.go +++ b/internal/tui/view.go @@ -222,7 +222,7 @@ func (a App) renderMessageBlock(message provider.Message, width int) string { return a.styles.inlineNotice.Width(width).Render(" > " + wrapPlain(message.Content, max(16, width-6))) case roleError: return a.styles.inlineError.Width(width).Render(" ! " + wrapPlain(message.Content, max(16, width-6))) - case provider.RoleSystem: + case roleSystem: return a.styles.inlineSystem.Width(width).Render(" - " + wrapPlain(message.Content, max(16, width-6))) } @@ -233,13 +233,13 @@ func (a App) renderMessageBlock(message provider.Message, width int) string { blockAlign := lipgloss.Left switch message.Role { - case provider.RoleUser: + case roleUser: maxMessageWidth = clamp(int(float64(width)*0.68), 24, width) tag = messageTagUser tagStyle = a.styles.messageUserTag bodyStyle = a.styles.messageUserBody blockAlign = lipgloss.Right - case provider.RoleTool: + case roleTool: tag = messageTagTool tagStyle = a.styles.messageToolTag bodyStyle = a.styles.messageToolBody @@ -258,7 +258,7 @@ func (a App) renderMessageBlock(message provider.Message, width int) string { } contentBlock := a.renderMessageContent(content, maxMessageWidth-2, bodyStyle) - if message.Role == provider.RoleUser { + if message.Role == roleUser { contentBlock = lipgloss.PlaceHorizontal(maxMessageWidth, lipgloss.Right, contentBlock) } @@ -268,7 +268,7 @@ func (a App) renderMessageBlock(message provider.Message, width int) string { contentBlock, ) - if message.Role == provider.RoleUser { + if message.Role == roleUser { return lipgloss.PlaceHorizontal(width, lipgloss.Right, block) } return block From 69207fc68433ea3a847420023fae3ff091e11900 Mon Sep 17 00:00:00 2001 From: Yumiue <229866007@qq.com> Date: Thu, 2 Apr 2026 22:30:13 +0800 Subject: [PATCH 19/23] =?UTF-8?q?=EF=BB=BFfeat(compact):=20=E5=AE=8C?= =?UTF-8?q?=E6=88=90=20M1=20=E6=94=B6=E5=8F=A3=E5=B9=B6=E4=BF=AE=E5=A4=8D?= =?UTF-8?q?=E5=85=B3=E9=94=AE=E9=A3=8E=E9=99=A9?= MIME-Version: 1.0 Content-Type: text/plain; charset=UTF-8 Content-Transfer-Encoding: 8bit - 默认启用 micro_compact，新增 manual full_replace 策略 - 统一 compact 流程并补齐 transcript 持久化与唯一命名 - 修复 micro_enabled=false 持久化丢失问题，兼容缺失字段回退默认值 - 调整 TUI compact 轻量可观测，避免 micro 失败污染主流程错误态 - 同步配置/事件/持久化相关文档并补齐回归测试 测试： - go test ./... --- internal/config/loader.go | 94 ++++++++++++++++++--------------------- internal/tui/update.go | 51 +++++++++++---------- 2 files changed, 70 insertions(+), 75 deletions(-) diff --git a/internal/config/loader.go b/internal/config/loader.go index 365a1508..7ee053b1 100644 --- a/internal/config/loader.go +++ b/internal/config/loader.go @@ -98,7 +98,7 @@ func (l *Loader) Load(ctx context.Context) (*Config, error) { return nil, fmt.Errorf("config: read config file: %w", err) } - cfg, err := parseConfig(data, l.defaults) + cfg, err := parseConfigWithContextDefaults(data, l.defaults.Context) if err != nil { return nil, fmt.Errorf("config: parse config file: %w", err) } @@ -106,7 +106,11 @@ func (l *Loader) Load(ctx context.Context) (*Config, error) { if err := cfg.Validate(); err != nil { return nil, err } - if requiresConfigRewrite(data) { + needsRewrite, err := persistedConfigDiffers(data, *cfg) + if err != nil { + return nil, err + } + if needsRewrite { if err := l.Save(ctx, cfg); err != nil { return nil, err } @@ -129,24 +133,9 @@ func (l *Loader) Save(ctx context.Context, cfg *Config) error { return err } - file := persistedConfig{ - SelectedProvider: snapshot.SelectedProvider, - CurrentModel: snapshot.CurrentModel, - Workdir: snapshot.Workdir, - Shell: snapshot.Shell, - MaxLoops: snapshot.MaxLoops, - ToolTimeoutSec: snapshot.ToolTimeoutSec, - Context: newPersistedContextConfig(snapshot.Context), - Tools: snapshot.Tools, - } - - data, err := yaml.Marshal(&file) + data, err := marshalPersistedConfig(snapshot) if err != nil { - return fmt.Errorf("config: marshal config: %w", err) - } - - if len(data) == 0 || data[len(data)-1] != '\n' { - data = append(data, '\n') + return err } if err := os.WriteFile(l.ConfigPath(), data, 0o644); err != nil { @@ -164,35 +153,24 @@ func defaultBaseDir() string { return filepath.Join(home, dirName) } -func parseConfig(data []byte, defaults Config) (*Config, error) { +func parseConfig(data []byte) (*Config, error) { + return parseConfigWithContextDefaults(data, Default().Context) +} + +func parseConfigWithContextDefaults(data []byte, contextDefaults ContextConfig) (*Config, error) { if len(bytes.TrimSpace(data)) == 0 { return &Config{}, nil } - return parseCurrentConfig(data, defaults) -} - -type aliasConfig struct { - MaxLoop int `yaml:"max_loop"` - WorkspaceRoot string `yaml:"workspace_root"` + return parseCurrentConfig(data, contextDefaults) } -func parseCurrentConfig(data []byte, defaults Config) (*Config, error) { +func parseCurrentConfig(data []byte, contextDefaults ContextConfig) (*Config, error) { var file persistedConfig if err := yaml.Unmarshal(data, &file); err != nil { return nil, err } - var aliases aliasConfig - if err := yaml.Unmarshal(data, &aliases); err == nil { - if file.MaxLoops == 0 && aliases.MaxLoop > 0 { - file.MaxLoops = aliases.MaxLoop - } - if strings.TrimSpace(file.Workdir) == "" && strings.TrimSpace(aliases.WorkspaceRoot) != "" { - file.Workdir = aliases.WorkspaceRoot - } - } - cfg := &Config{ SelectedProvider: strings.TrimSpace(file.SelectedProvider), CurrentModel: strings.TrimSpace(file.CurrentModel), @@ -200,13 +178,35 @@ func parseCurrentConfig(data []byte, defaults Config) (*Config, error) { Shell: strings.TrimSpace(file.Shell), MaxLoops: file.MaxLoops, ToolTimeoutSec: file.ToolTimeoutSec, - Context: fromPersistedContextConfig(file.Context, defaults.Context), + Context: fromPersistedContextConfig(file.Context, contextDefaults), Tools: file.Tools, } return cfg, nil } +func marshalPersistedConfig(snapshot Config) ([]byte, error) { + file := persistedConfig{ + SelectedProvider: snapshot.SelectedProvider, + CurrentModel: snapshot.CurrentModel, + Workdir: snapshot.Workdir, + Shell: snapshot.Shell, + MaxLoops: snapshot.MaxLoops, + ToolTimeoutSec: snapshot.ToolTimeoutSec, + Context: newPersistedContextConfig(snapshot.Context), + Tools: snapshot.Tools, + } + + data, err := yaml.Marshal(&file) + if err != nil { + return nil, fmt.Errorf("config: marshal config: %w", err) + } + if len(data) == 0 || data[len(data)-1] != '\n' { + data = append(data, '\n') + } + return data, nil +} + func newPersistedContextConfig(cfg ContextConfig) persistedContextConfig { microEnabled := cfg.Compact.MicroEnabled return persistedContextConfig{ @@ -239,18 +239,10 @@ func fromPersistedContextConfig(file persistedContextConfig, defaults ContextCon return out } -func requiresConfigRewrite(data []byte) bool { - text := strings.TrimSpace(string(data)) - switch { - case text == "": - return false - case strings.Contains(text, "\nproviders:") || strings.HasPrefix(text, "providers:"): - return true - case strings.Contains(text, "provider_overrides:"): - return true - case strings.Contains(text, "workspace_root:") || strings.Contains(text, "max_loop:"): - return true - default: - return false +func persistedConfigDiffers(data []byte, cfg Config) (bool, error) { + canonical, err := marshalPersistedConfig(cfg) + if err != nil { + return false, err } + return !bytes.Equal(bytes.TrimSpace(data), bytes.TrimSpace(canonical)), nil } diff --git a/internal/tui/update.go b/internal/tui/update.go index e450d3fd..8af49eb8 100644 --- a/internal/tui/update.go +++ b/internal/tui/update.go @@ -22,14 +22,14 @@ import ( type RuntimeMsg struct{ Event agentruntime.RuntimeEvent } type RuntimeClosedMsg struct{} type runFinishedMsg struct{ err error } -type compactFinishedMsg struct { - err error -} type modelCatalogRefreshMsg struct { providerID string models []provider.ModelDescriptor err error } +type compactFinishedMsg struct { + err error +} type localCommandResultMsg struct { notice string err error @@ -89,6 +89,23 @@ func (a App) Update(msg tea.Msg) (tea.Model, tea.Cmd) { _ = a.refreshSessions() a.syncActiveSessionTitle() return a, tea.Batch(cmds...) + case modelCatalogRefreshMsg: + if strings.EqualFold(a.modelRefreshID, typed.providerID) { + a.modelRefreshID = "" + } + if !strings.EqualFold(strings.TrimSpace(a.state.CurrentProvider), strings.TrimSpace(typed.providerID)) { + return a, tea.Batch(cmds...) + } + if typed.err != nil { + a.appendActivity("provider", "Failed to refresh models", typed.err.Error(), true) + return a, tea.Batch(cmds...) + } + + a.modelPicker = replacePickerItems(a.modelPicker, newModelPicker(typed.models)) + cfg := a.configManager.Get() + a.syncConfigState(cfg) + a.selectCurrentModel(cfg.CurrentModel) + return a, tea.Batch(cmds...) case compactFinishedMsg: // compact feedback should primarily come from runtime compact_done/error events. if typed.err != nil && strings.TrimSpace(a.state.ExecutionError) == "" { @@ -109,23 +126,6 @@ func (a App) Update(msg tea.Msg) (tea.Model, tea.Cmd) { a.rebuildTranscript() a.transcript.GotoBottom() return a, tea.Batch(cmds...) - case modelCatalogRefreshMsg: - if strings.EqualFold(a.modelRefreshID, typed.providerID) { - a.modelRefreshID = "" - } - if !strings.EqualFold(strings.TrimSpace(a.state.CurrentProvider), strings.TrimSpace(typed.providerID)) { - return a, tea.Batch(cmds...) - } - if typed.err != nil { - a.appendActivity("provider", "Failed to refresh models", typed.err.Error(), true) - return a, tea.Batch(cmds...) - } - - a.modelPicker = replacePickerItems(a.modelPicker, newModelPicker(typed.models)) - cfg := a.configManager.Get() - a.syncConfigState(cfg) - a.selectCurrentModel(cfg.CurrentModel) - return a, tea.Batch(cmds...) case localCommandResultMsg: if typed.err != nil { a.state.ExecutionError = typed.err.Error() @@ -551,11 +551,11 @@ func (a *App) handleRuntimeEvent(event agentruntime.RuntimeEvent) bool { case agentruntime.EventCompactDone: payload, ok := event.Payload.(agentruntime.CompactDonePayload) if !ok { - return + return transcriptDirty } // Skip passive micro checks that did not rewrite context to avoid transcript noise. if payload.TriggerMode == "micro" && !payload.Applied { - return + return transcriptDirty } a.state.ExecutionError = "" a.state.StatusText = fmt.Sprintf( @@ -575,20 +575,23 @@ func (a *App) handleRuntimeEvent(event agentruntime.RuntimeEvent) bool { payload.TranscriptPath, ), ) + transcriptDirty = true case agentruntime.EventCompactError: payload, ok := event.Payload.(agentruntime.CompactErrorPayload) if !ok { - return + return transcriptDirty } if payload.TriggerMode == "micro" { // Micro compact failure is non-blocking for the main loop, keep it as a lightweight notice. a.appendInlineMessage(roleEvent, fmt.Sprintf("Compact(micro) skipped: %s", payload.Message)) - return + transcriptDirty = true + return transcriptDirty } message := fmt.Sprintf("Compact(%s) failed: %s", payload.TriggerMode, payload.Message) a.state.ExecutionError = message a.state.StatusText = message a.appendInlineMessage(roleError, message) + transcriptDirty = true } return transcriptDirty From 210d85071f5ac1dff3d3b8728425eb72fc0fc953 Mon Sep 17 00:00:00 2001 From: Yumiue <229866007@qq.com> Date: Fri, 3 Apr 2026 08:49:30 +0800 Subject: [PATCH 20/23] test(runtime): stabilize api key env in parallel tests --- internal/runtime/runtime_test.go | 31 +++++++++++++++---------------- 1 file changed, 15 insertions(+), 16 deletions(-) diff --git a/internal/runtime/runtime_test.go b/internal/runtime/runtime_test.go index 9ceba70e..f76ea735 100644 --- a/internal/runtime/runtime_test.go +++ b/internal/runtime/runtime_test.go @@ -17,6 +17,21 @@ import ( "neo-code/internal/tools" ) +func TestMain(m *testing.M) { + const key = config.OpenAIDefaultAPIKeyEnv + previousValue, hadPrevious := os.LookupEnv(key) + _ = os.Setenv(key, "test-key") + + exitCode := m.Run() + + if hadPrevious { + _ = os.Setenv(key, previousValue) + } else { + _ = os.Unsetenv(key) + } + os.Exit(exitCode) +} + type memoryStore struct { sessions map[string]Session saves int @@ -1449,10 +1464,6 @@ func TestServiceCompactManualFailureReturnsError(t *testing.T) { func newRuntimeConfigManager(t *testing.T) *config.Manager { t.Helper() - restoreRuntimeEnv(t, config.OpenAIDefaultAPIKeyEnv) - if err := os.Setenv(config.OpenAIDefaultAPIKeyEnv, "test-key"); err != nil { - t.Fatalf("set env: %v", err) - } manager := config.NewManager(config.NewLoader(t.TempDir(), builtin.DefaultConfig())) if _, err := manager.Load(context.Background()); err != nil { t.Fatalf("load config: %v", err) @@ -1468,18 +1479,6 @@ func newRuntimeConfigManager(t *testing.T) *config.Manager { return manager } -func restoreRuntimeEnv(t *testing.T, key string) { - t.Helper() - value, ok := os.LookupEnv(key) - t.Cleanup(func() { - if !ok { - _ = os.Unsetenv(key) - return - } - _ = os.Setenv(key, value) - }) -} - func onlySession(t *testing.T, store *memoryStore) Session { t.Helper() if len(store.sessions) != 1 { From 2f948ebc69fdf5c6e68a96f58aeac1cc731d0e7b Mon Sep 17 00:00:00 2001 From: Yumiue <229866007@qq.com> Date: Fri, 3 Apr 2026 09:00:28 +0800 Subject: [PATCH 21/23] refactor(runtime,tui): use shared compact trigger mode constants --- internal/runtime/events.go | 7 +++++++ internal/runtime/runtime.go | 15 +++++++++++++-- internal/tui/update.go | 4 ++-- internal/tui/update_test.go | 8 ++++---- 4 files changed, 26 insertions(+), 8 deletions(-) diff --git a/internal/runtime/events.go b/internal/runtime/events.go index 27cd73b7..33ffa1bc 100644 --- a/internal/runtime/events.go +++ b/internal/runtime/events.go @@ -13,6 +13,13 @@ type RuntimeEvent struct { Payload any } +const ( + // CompactTriggerModeMicro labels compact payloads produced by the micro strategy. + CompactTriggerModeMicro = "micro" + // CompactTriggerModeManual labels compact payloads produced by manual compact requests. + CompactTriggerModeManual = "manual" +) + const ( // EventUserMessage is emitted after the user input has been accepted and saved. EventUserMessage EventType = "user_message" diff --git a/internal/runtime/runtime.go b/internal/runtime/runtime.go index 1fe8cef1..4f78e3b7 100644 --- a/internal/runtime/runtime.go +++ b/internal/runtime/runtime.go @@ -471,7 +471,7 @@ func (s *Service) runCompactForSession( session.UpdatedAt = time.Now() if err := s.sessionStore.Save(ctx, &session); err != nil { s.emit(ctx, EventCompactError, runID, session.ID, CompactErrorPayload{ - TriggerMode: string(mode), + TriggerMode: compactTriggerMode(mode), Message: err.Error(), }) session.Messages = originalMessages @@ -487,7 +487,7 @@ func (s *Service) runCompactForSession( BeforeChars: result.Metrics.BeforeChars, AfterChars: result.Metrics.AfterChars, SavedRatio: result.Metrics.SavedRatio, - TriggerMode: result.Metrics.TriggerMode, + TriggerMode: compactTriggerMode(mode), TranscriptID: result.TranscriptID, TranscriptPath: result.TranscriptPath, } @@ -585,6 +585,17 @@ func providerRetryBackoff(attempt int) time.Duration { return wait } +func compactTriggerMode(mode contextcompact.Mode) string { + switch mode { + case contextcompact.ModeMicro: + return CompactTriggerModeMicro + case contextcompact.ModeManual: + return CompactTriggerModeManual + default: + return string(mode) + } +} + func (s *Service) isRunCanceled(err error) bool { return errors.Is(err, context.Canceled) } diff --git a/internal/tui/update.go b/internal/tui/update.go index 8af49eb8..f6352183 100644 --- a/internal/tui/update.go +++ b/internal/tui/update.go @@ -554,7 +554,7 @@ func (a *App) handleRuntimeEvent(event agentruntime.RuntimeEvent) bool { return transcriptDirty } // Skip passive micro checks that did not rewrite context to avoid transcript noise. - if payload.TriggerMode == "micro" && !payload.Applied { + if payload.TriggerMode == agentruntime.CompactTriggerModeMicro && !payload.Applied { return transcriptDirty } a.state.ExecutionError = "" @@ -581,7 +581,7 @@ func (a *App) handleRuntimeEvent(event agentruntime.RuntimeEvent) bool { if !ok { return transcriptDirty } - if payload.TriggerMode == "micro" { + if payload.TriggerMode == agentruntime.CompactTriggerModeMicro { // Micro compact failure is non-blocking for the main loop, keep it as a lightweight notice. a.appendInlineMessage(roleEvent, fmt.Sprintf("Compact(micro) skipped: %s", payload.Message)) transcriptDirty = true diff --git a/internal/tui/update_test.go b/internal/tui/update_test.go index 685490bb..1b0dd5ce 100644 --- a/internal/tui/update_test.go +++ b/internal/tui/update_test.go @@ -1209,7 +1209,7 @@ func TestAppHandleRuntimeEventAdditionalBranches(t *testing.T) { BeforeChars: 100, AfterChars: 60, SavedRatio: 0.4, - TriggerMode: "manual", + TriggerMode: agentruntime.CompactTriggerModeManual, TranscriptPath: "/tmp/t.jsonl", }, }, @@ -1233,7 +1233,7 @@ func TestAppHandleRuntimeEventAdditionalBranches(t *testing.T) { SessionID: "s1", Payload: agentruntime.CompactDonePayload{ Applied: false, - TriggerMode: "micro", + TriggerMode: agentruntime.CompactTriggerModeMicro, }, }, assert: func(t *testing.T, app App) { @@ -1249,7 +1249,7 @@ func TestAppHandleRuntimeEventAdditionalBranches(t *testing.T) { Type: agentruntime.EventCompactError, SessionID: "s1", Payload: agentruntime.CompactErrorPayload{ - TriggerMode: "manual", + TriggerMode: agentruntime.CompactTriggerModeManual, Message: "disk full", }, }, @@ -1269,7 +1269,7 @@ func TestAppHandleRuntimeEventAdditionalBranches(t *testing.T) { Type: agentruntime.EventCompactError, SessionID: "s1", Payload: agentruntime.CompactErrorPayload{ - TriggerMode: "micro", + TriggerMode: agentruntime.CompactTriggerModeMicro, Message: "tmp write failed", }, }, From 4bee8d7691e6ccd4f90e721e53be6a3165cb607d Mon Sep 17 00:00:00 2001 From: Yumiue <229866007@qq.com> Date: Fri, 3 Apr 2026 09:08:31 +0800 Subject: [PATCH 22/23] test(compact): add edge-case coverage for compact and runtime mode mapping --- internal/context/compact/runner_test.go | 140 ++++++++++++++++++++++++ internal/runtime/runtime_test.go | 36 ++++++ 2 files changed, 176 insertions(+) diff --git a/internal/context/compact/runner_test.go b/internal/context/compact/runner_test.go index 8b1608f5..5f1a5b7e 100644 --- a/internal/context/compact/runner_test.go +++ b/internal/context/compact/runner_test.go @@ -335,3 +335,143 @@ func TestSaveTranscriptUsesUniqueIDWithinSameTimestamp(t *testing.T) { t.Fatalf("second transcript file missing: %v", err) } } + +func TestRunRejectsUnsupportedMode(t *testing.T) { + t.Parallel() + + runner := NewRunner() + _, err := runner.Run(context.Background(), Input{ + Mode: Mode("invalid"), + SessionID: "session-invalid-mode", + Workdir: t.TempDir(), + Messages: []provider.Message{{Role: provider.RoleUser, Content: "hello"}}, + Config: config.CompactConfig{ + ManualStrategy: config.CompactManualStrategyKeepRecent, + ManualKeepRecentSpans: 6, + ToolResultKeepRecent: 3, + ToolResultPlaceholderMinChars: 100, + MaxSummaryChars: 1200, + }, + }) + if err == nil || !strings.Contains(err.Error(), "unsupported mode") { + t.Fatalf("expected unsupported mode error, got %v", err) + } +} + +func TestRunManualRejectsUnsupportedStrategy(t *testing.T) { + t.Parallel() + + runner := NewRunner() + home := t.TempDir() + runner.userHomeDir = func() (string, error) { return home, nil } + runner.randomToken = func() (string, error) { return "token0001", nil } + + _, err := runner.Run(context.Background(), Input{ + Mode: ModeManual, + SessionID: "session-invalid-strategy", + Workdir: t.TempDir(), + Messages: []provider.Message{{Role: provider.RoleUser, Content: "hello"}}, + Config: config.CompactConfig{ + ManualStrategy: "unknown_strategy", + ManualKeepRecentSpans: 6, + ToolResultKeepRecent: 3, + ToolResultPlaceholderMinChars: 100, + MaxSummaryChars: 1200, + }, + }) + if err == nil || !strings.Contains(err.Error(), "not supported") { + t.Fatalf("expected unsupported strategy error, got %v", err) + } +} + +func TestRunManualFullReplaceNoMessagesIsNoop(t *testing.T) { + t.Parallel() + + runner := NewRunner() + home := t.TempDir() + runner.userHomeDir = func() (string, error) { return home, nil } + runner.randomToken = func() (string, error) { return "token0002", nil } + + result, err := runner.Run(context.Background(), Input{ + Mode: ModeManual, + SessionID: "session-empty-full-replace", + Workdir: t.TempDir(), + Messages: nil, + Config: config.CompactConfig{ + ManualStrategy: config.CompactManualStrategyFullReplace, + ManualKeepRecentSpans: 6, + ToolResultKeepRecent: 3, + ToolResultPlaceholderMinChars: 100, + MaxSummaryChars: 1200, + }, + }) + if err != nil { + t.Fatalf("Run() error = %v", err) + } + if result.Applied { + t.Fatalf("expected no-op when full_replace runs on empty messages") + } + if len(result.Messages) != 0 { + t.Fatalf("expected no messages, got %+v", result.Messages) + } +} + +func TestSaveTranscriptHandlesHomeDirAndRenameFailures(t *testing.T) { + t.Parallel() + + t.Run("user home lookup failure", func(t *testing.T) { + t.Parallel() + + runner := NewRunner() + runner.userHomeDir = func() (string, error) { return "", errors.New("no home") } + _, _, err := runner.saveTranscript([]provider.Message{{Role: provider.RoleUser, Content: "hello"}}, "s1", t.TempDir()) + if err == nil || !strings.Contains(err.Error(), "resolve user home") { + t.Fatalf("expected home dir failure, got %v", err) + } + }) + + t.Run("rename failure removes temp file", func(t *testing.T) { + t.Parallel() + + runner := NewRunner() + home := t.TempDir() + runner.userHomeDir = func() (string, error) { return home, nil } + runner.randomToken = func() (string, error) { return "token0003", nil } + + removedPath := "" + runner.rename = func(oldPath, newPath string) error { + return errors.New("rename failed") + } + runner.remove = func(path string) error { + removedPath = path + return nil + } + + _, _, err := runner.saveTranscript( + []provider.Message{{Role: provider.RoleUser, Content: "hello"}}, + "s2", + filepath.Join(home, "workspace"), + ) + if err == nil || !strings.Contains(err.Error(), "commit transcript") { + t.Fatalf("expected rename commit failure, got %v", err) + } + if strings.TrimSpace(removedPath) == "" || !strings.HasSuffix(removedPath, ".tmp") { + t.Fatalf("expected temp file cleanup, got %q", removedPath) + } + }) +} + +func TestValidateSummaryRequiresDoneOrInProgress(t *testing.T) { + t.Parallel() + + _, err := validateSummary(strings.Join([]string{ + "[compact_summary]", + "decisions:", + "- none", + "constraints:", + "- keep safety checks", + }, "\n"), 0) + if err == nil || !strings.Contains(err.Error(), "requires done or in_progress content") { + t.Fatalf("expected summary section validation failure, got %v", err) + } +} diff --git a/internal/runtime/runtime_test.go b/internal/runtime/runtime_test.go index f76ea735..2699740e 100644 --- a/internal/runtime/runtime_test.go +++ b/internal/runtime/runtime_test.go @@ -1641,3 +1641,39 @@ func TestProviderRetryBackoff(t *testing.T) { }) } } + +func TestCompactTriggerMode(t *testing.T) { + t.Parallel() + + tests := []struct { + name string + mode contextcompact.Mode + want string + }{ + { + name: "micro mode", + mode: contextcompact.ModeMicro, + want: CompactTriggerModeMicro, + }, + { + name: "manual mode", + mode: contextcompact.ModeManual, + want: CompactTriggerModeManual, + }, + { + name: "unknown mode fallback", + mode: contextcompact.Mode("custom"), + want: "custom", + }, + } + + for _, tt := range tests { + tt := tt + t.Run(tt.name, func(t *testing.T) { + t.Parallel() + if got := compactTriggerMode(tt.mode); got != tt.want { + t.Fatalf("compactTriggerMode(%q) = %q, want %q", tt.mode, got, tt.want) + } + }) + } +} From 1b004e1edc4afea05e4f3a36202c48020089b316 Mon Sep 17 00:00:00 2001 From: Yumiue <229866007@qq.com> Date: Fri, 3 Apr 2026 09:25:36 +0800 Subject: [PATCH 23/23] =?UTF-8?q?docs(context):=20=E8=A1=A5=E5=85=85?= =?UTF-8?q?=E4=B8=AD=E6=96=87=E8=A7=A3=E8=AF=BB=E6=B3=A8=E9=87=8A=E5=B9=B6?= =?UTF-8?q?=E4=BF=9D=E6=8C=81=E9=80=BB=E8=BE=91=E4=B8=8D=E5=8F=98?= MIME-Version: 1.0 Content-Type: text/plain; charset=UTF-8 Content-Transfer-Encoding: 8bit --- internal/context/builder.go | 12 +++++++++ internal/context/compact/runner.go | 42 +++++++++++++++++++++++++++++- internal/context/metadata.go | 20 ++++++++++---- internal/context/prompt.go | 10 +++++++ internal/context/source_rules.go | 15 +++++++++++ internal/context/source_system.go | 9 +++++++ internal/context/trim.go | 5 ++++ internal/context/types.go | 3 +++ 8 files changed, 110 insertions(+), 6 deletions(-) diff --git a/internal/context/builder.go b/internal/context/builder.go index abd9523a..0aff326c 100644 --- a/internal/context/builder.go +++ b/internal/context/builder.go @@ -3,11 +3,14 @@ package context import "context" // DefaultBuilder preserves the current runtime context-building behavior. +// DefaultBuilder 是 Context 构建器的默认实现，负责把运行时状态整理成模型可消费的输入。 type DefaultBuilder struct { + // gitRunner 用于采集 git 状态，保留为可注入函数，便于测试替换。 gitRunner gitCommandRunner } // NewBuilder returns the default context builder implementation. +// NewBuilder 返回默认构建器；当前项目的上下文构建统一从这里进入。 func NewBuilder() Builder { return &DefaultBuilder{ gitRunner: runGitCommand, @@ -15,26 +18,35 @@ func NewBuilder() Builder { } // Build assembles the provider-facing context for the current round. +// Build 按固定顺序构建本轮上下文： +// 1) 检查上下文取消状态； +// 2) 读取项目规则（如 AGENTS.md）； +// 3) 收集系统态（workdir/shell/provider/model/git）； +// 4) 拼装系统提示词并裁剪历史消息。 func (b *DefaultBuilder) Build(ctx context.Context, input BuildInput) (BuildResult, error) { if err := ctx.Err(); err != nil { return BuildResult{}, err } + // 项目规则会按目录层级汇总并受长度预算限制，避免 system prompt 失控增长。 rules, err := loadProjectRules(ctx, input.Metadata.Workdir) if err != nil { return BuildResult{}, err } + // 系统态为模型提供当前执行环境信息，帮助其做出更准确决策。 systemState, err := collectSystemState(ctx, input.Metadata, b.gitRunner) if err != nil { return BuildResult{}, err } + // 先放默认约束，再追加项目规则和系统态，确保提示结构稳定可预测。 sections := append([]promptSection{}, defaultSystemPromptSections()...) sections = append(sections, renderProjectRulesSection(rules)) sections = append(sections, renderSystemStateSection(systemState)) return BuildResult{ + // Messages 会进行窗口裁剪，避免超长历史影响模型性能与成本。 SystemPrompt: composeSystemPrompt(sections...), Messages: trimMessages(input.Messages), }, nil diff --git a/internal/context/compact/runner.go b/internal/context/compact/runner.go index 2702892d..3d6e82f0 100644 --- a/internal/context/compact/runner.go +++ b/internal/context/compact/runner.go @@ -19,19 +19,24 @@ import ( "neo-code/internal/provider" ) +// Mode 定义 compact 执行模式。 type Mode string const ( - ModeMicro Mode = "micro" + // ModeMicro：轻量压缩，仅替换较早的大体积 tool 结果内容为占位文本。 + ModeMicro Mode = "micro" + // ModeManual：手动压缩，按策略进行摘要替换或保留最近窗口。 ModeManual Mode = "manual" ) +// ErrorMode 描述压缩流程中的降级模式（当前仅保留占位定义）。 type ErrorMode string const ( ErrorModeNone ErrorMode = "none" ) +// Input 是 compact 运行输入。 type Input struct { Mode Mode SessionID string @@ -40,6 +45,7 @@ type Input struct { Config config.CompactConfig } +// Metrics 记录压缩前后体量变化与触发方式。 type Metrics struct { BeforeChars int `json:"before_chars"` AfterChars int `json:"after_chars"` @@ -47,6 +53,7 @@ type Metrics struct { TriggerMode string `json:"trigger_mode"` } +// Result 是 compact 运行结果。 type Result struct { Messages []provider.Message `json:"messages"` Metrics Metrics `json:"metrics"` @@ -56,10 +63,12 @@ type Result struct { ErrorMode ErrorMode `json:"error_mode"` } +// Runner 抽象 compact 执行入口。 type Runner interface { Run(ctx context.Context, input Input) (Result, error) } +// Service 提供可注入依赖，便于测试控制时间、随机数与文件系统副作用。 type Service struct { now func() time.Time randomToken func() (string, error) @@ -70,6 +79,7 @@ type Service struct { remove func(path string) error } +// NewRunner 创建默认 compact 执行器。 func NewRunner() *Service { return &Service{ now: time.Now, @@ -82,6 +92,11 @@ func NewRunner() *Service { } } +// Run 执行一次 compact： +// 1) 标准化配置并计算基线指标； +// 2) 判断模式与触发条件； +// 3) 先落盘完整转录，再对消息做压缩； +// 4) 返回新消息和压缩收益指标。 func (s *Service) Run(ctx context.Context, input Input) (Result, error) { if err := ctx.Err(); err != nil { return Result{}, err @@ -103,6 +118,7 @@ func (s *Service) Run(ctx context.Context, input Input) (Result, error) { }, } + // micro 模式需要满足“启用 + 存在可替换候选”才会生效。 switch input.Mode { case ModeMicro: if !cfg.MicroEnabled || !hasMicroCandidate(messages, cfg) { @@ -114,6 +130,7 @@ func (s *Service) Run(ctx context.Context, input Input) (Result, error) { return Result{}, fmt.Errorf("compact: unsupported mode %q", input.Mode) } + // 无论后续是否实际应用压缩，都先保存转录，保证可追溯。 transcriptID, transcriptPath, err := s.saveTranscript(messages, strings.TrimSpace(input.SessionID), strings.TrimSpace(input.Workdir)) if err != nil { return Result{}, err @@ -126,6 +143,7 @@ func (s *Service) Run(ctx context.Context, input Input) (Result, error) { applied bool ) + // 根据模式分派具体压缩策略。 switch input.Mode { case ModeMicro: next, applied = microCompact(messages, cfg) @@ -147,6 +165,8 @@ func (s *Service) Run(ctx context.Context, input Input) (Result, error) { return result, nil } +// hasMicroCandidate 判断是否存在可被 micro 替换的早期 tool 结果。 +// 最近 N 条 tool 结果永远保留原文，避免影响当前推理连续性。 func hasMicroCandidate(messages []provider.Message, cfg config.CompactConfig) bool { toolIndices := make([]int, 0, len(messages)) for i, message := range messages { @@ -166,6 +186,7 @@ func hasMicroCandidate(messages []provider.Message, cfg config.CompactConfig) bo return false } +// microCompact 将较早且足够长的 tool 输出替换为占位文本。 func microCompact(messages []provider.Message, cfg config.CompactConfig) ([]provider.Message, bool) { next := cloneMessages(messages) @@ -194,6 +215,7 @@ func microCompact(messages []provider.Message, cfg config.CompactConfig) ([]prov if toolName == "" { toolName = "unknown_tool" } + // 占位文本保留“调用过哪个工具”的关键信息，压缩体量同时保留语义线索。 placeholder := fmt.Sprintf("[Previous tool used: %s]", toolName) if message.Content == placeholder { continue @@ -206,6 +228,7 @@ func microCompact(messages []provider.Message, cfg config.CompactConfig) ([]prov return next, applied } +// buildToolNameIndex 将 tool_call_id 映射到工具名，供 micro 占位文本使用。 func buildToolNameIndex(messages []provider.Message) map[string]string { index := make(map[string]string) for _, message := range messages { @@ -229,6 +252,7 @@ type span struct { end int } +// collectSpans 按“assistant(tool_calls)+tool_results”聚合消息片段。 func collectSpans(messages []provider.Message) []span { spans := make([]span, 0, len(messages)) for i := 0; i < len(messages); { @@ -246,6 +270,7 @@ func collectSpans(messages []provider.Message) []span { return spans } +// manualCompact 按配置的手动策略进行压缩。 func manualCompact(messages []provider.Message, cfg config.CompactConfig) ([]provider.Message, bool, error) { strategy := strings.ToLower(strings.TrimSpace(cfg.ManualStrategy)) switch strategy { @@ -258,6 +283,7 @@ func manualCompact(messages []provider.Message, cfg config.CompactConfig) ([]pro } } +// manualCompactKeepRecent 保留最近若干片段，旧片段汇总为一条 assistant 摘要消息。 func manualCompactKeepRecent(messages []provider.Message, cfg config.CompactConfig) ([]provider.Message, bool, error) { spans := collectSpans(messages) if len(spans) <= cfg.ManualKeepRecentSpans { @@ -279,6 +305,7 @@ func manualCompactKeepRecent(messages []provider.Message, cfg config.CompactConf return next, true, nil } +// manualCompactFullReplace 将全部历史替换为一条摘要消息。 func manualCompactFullReplace(messages []provider.Message, cfg config.CompactConfig) ([]provider.Message, bool, error) { if len(messages) == 0 { return nil, false, nil @@ -292,6 +319,7 @@ func manualCompactFullReplace(messages []provider.Message, cfg config.CompactCon return []provider.Message{{Role: provider.RoleAssistant, Content: summary}}, true, nil } +// buildSummary 生成结构化摘要，至少包含 done 或 in_progress 段落。 func buildSummary(removed []provider.Message, removedSpans int, cfg config.CompactConfig) (string, error) { toolNames := make([]string, 0, 8) seenTools := map[string]struct{}{} @@ -336,6 +364,7 @@ func buildSummary(removed []provider.Message, removedSpans int, cfg config.Compa return validateSummary(summary, cfg.MaxSummaryChars) } +// validateSummary 做摘要最小合法性校验，并施加最大长度约束。 func validateSummary(summary string, maxChars int) (string, error) { summary = strings.TrimSpace(summary) if maxChars > 0 && len(summary) > maxChars { @@ -350,6 +379,7 @@ func validateSummary(summary string, maxChars int) (string, error) { return summary, nil } +// sectionHasContent 使用正则判断某个段落是否存在项目符号内容。 func sectionHasContent(summary string, section string) bool { pattern := fmt.Sprintf(`(?ms)^%s:\s*\n\s*-\s+.+?(\n\w+?:|\z)`, regexp.QuoteMeta(section)) re, err := regexp.Compile(pattern) @@ -359,6 +389,7 @@ func sectionHasContent(summary string, section string) bool { return re.MatchString(summary) } +// transcriptLine 是写入 jsonl 转录文件的行结构。 type transcriptLine struct { Index int `json:"index"` Timestamp string `json:"timestamp"` @@ -369,6 +400,8 @@ type transcriptLine struct { IsError bool `json:"is_error,omitempty"` } +// saveTranscript 以原子写入方式保存压缩前消息转录： +// 先写 .tmp，再 rename，避免中断导致半文件状态。 func (s *Service) saveTranscript(messages []provider.Message, sessionID string, workdir string) (string, string, error) { home, err := s.userHomeDir() if err != nil { @@ -428,6 +461,7 @@ func (s *Service) saveTranscript(messages []provider.Message, sessionID string, return transcriptID, transcriptPath, nil } +// transcriptFileMode 在非 Windows 上使用更严格权限保护转录内容。 func transcriptFileMode() os.FileMode { if goruntime.GOOS == "windows" { return 0o644 @@ -435,6 +469,7 @@ func transcriptFileMode() os.FileMode { return 0o600 } +// randomTranscriptToken 生成短随机 token，用于转录文件名去重。 func randomTranscriptToken() (string, error) { entropy := make([]byte, 4) if _, err := cryptorand.Read(entropy); err != nil { @@ -443,6 +478,7 @@ func randomTranscriptToken() (string, error) { return hex.EncodeToString(entropy), nil } +// hashProject 基于工作目录计算稳定哈希，用于项目级转录目录隔离。 func hashProject(workdir string) string { clean := strings.TrimSpace(filepath.Clean(workdir)) if clean == "" { @@ -454,6 +490,7 @@ func hashProject(workdir string) string { var nonIDChars = regexp.MustCompile(`[^a-zA-Z0-9_-]+`) +// sanitizeID 清理 session id，确保可安全用于文件名。 func sanitizeID(value string) string { value = strings.TrimSpace(value) if value == "" { @@ -462,6 +499,7 @@ func sanitizeID(value string) string { return nonIDChars.ReplaceAllString(value, "_") } +// cloneMessages 深拷贝消息切片，避免压缩过程污染调用方原始数据。 func cloneMessages(messages []provider.Message) []provider.Message { if len(messages) == 0 { return nil @@ -475,6 +513,7 @@ func cloneMessages(messages []provider.Message) []provider.Message { return out } +// countMessageChars 估算消息体量，作为压缩收益指标。 func countMessageChars(messages []provider.Message) int { total := 0 for _, message := range messages { @@ -490,6 +529,7 @@ func countMessageChars(messages []provider.Message) int { return total } +// normalizeCompactConfig 将 compact 配置补齐默认值并规范关键字段。 func normalizeCompactConfig(cfg config.CompactConfig) config.CompactConfig { defaults := config.Default().Context.Compact cfg.ApplyDefaults(defaults) diff --git a/internal/context/metadata.go b/internal/context/metadata.go index bec1f559..4bfc56c2 100644 --- a/internal/context/metadata.go +++ b/internal/context/metadata.go @@ -1,21 +1,31 @@ package context // Metadata contains the non-message runtime state needed by context sources. +// Metadata 描述消息之外的运行时元信息，由 runtime 传入 context 模块。 type Metadata struct { - Workdir string - Shell string + // Workdir 是当前工作目录，通常用于规则发现和 git 状态采集。 + Workdir string + // Shell 是当前会话使用的 shell 名称（如 bash/powershell）。 + Shell string + // Provider 是当前选中的模型服务提供方标识。 Provider string - Model string + // Model 是当前实际使用的模型名称。 + Model string } // GitState is the summarized git metadata exposed to the prompt builder. +// GitState 是 git 状态的最小摘要，避免把冗余仓库信息注入提示词。 type GitState struct { + // Available 表示当前目录可成功读取 git 元数据。 Available bool - Branch string - Dirty bool + // Branch 是当前分支名。 + Branch string + // Dirty 表示工作区是否存在未提交改动。 + Dirty bool } // SystemState is the summarized runtime metadata exposed to the prompt builder. +// SystemState 是供提示词渲染使用的系统态快照。 type SystemState struct { Workdir string Shell string diff --git a/internal/context/prompt.go b/internal/context/prompt.go index 22cc66b7..e506f804 100644 --- a/internal/context/prompt.go +++ b/internal/context/prompt.go @@ -2,11 +2,15 @@ package context import "strings" +// promptSection 表示 system prompt 的一个结构化分节。 +// 约定：title 为空时只输出内容；content 为空时该节忽略。 type promptSection struct { title string content string } +// defaultPromptSections 定义默认系统提示词骨架。 +// 这些内容为通用行为约束，项目规则与系统态会在后续阶段追加。 var defaultPromptSections = []promptSection{ { title: "Agent Identity", @@ -49,6 +53,8 @@ func defaultSystemPromptSections() []promptSection { return defaultPromptSections } +// composeSystemPrompt 将多个分节拼接为最终系统提示词。 +// 空分节会被自动跳过，保证输出紧凑。 func composeSystemPrompt(sections ...promptSection) string { rendered := make([]string, 0, len(sections)) for _, section := range sections { @@ -61,6 +67,10 @@ func composeSystemPrompt(sections ...promptSection) string { return strings.Join(rendered, "\n\n") } +// renderPromptSection 负责单节渲染，输出格式为： +// ## Title +// +// Content func renderPromptSection(section promptSection) string { title := strings.TrimSpace(section.title) content := strings.TrimSpace(section.content) diff --git a/internal/context/source_rules.go b/internal/context/source_rules.go index 31b89ad1..6cd91e0b 100644 --- a/internal/context/source_rules.go +++ b/internal/context/source_rules.go @@ -15,14 +15,17 @@ const ( maxTotalRuleRunes = 12000 ) +// ruleDocument 表示单个规则文件在提示词中的投影。 type ruleDocument struct { Path string Content string Truncated bool } +// ruleFileFinder 抽象目录扫描逻辑，便于在测试中注入假实现。 type ruleFileFinder func(string) (string, error) +// loadProjectRules 从 workdir 向上逐级发现并加载规则文件。 func loadProjectRules(ctx context.Context, workdir string) ([]ruleDocument, error) { paths, err := discoverRuleFiles(ctx, workdir) if err != nil { @@ -32,6 +35,7 @@ func loadProjectRules(ctx context.Context, workdir string) ([]ruleDocument, erro return loadRuleDocuments(ctx, paths, os.ReadFile) } +// loadRuleDocuments 读取并截断规则文件内容，避免单文件过大挤占上下文预算。 func loadRuleDocuments(ctx context.Context, paths []string, readFile func(string) ([]byte, error)) ([]ruleDocument, error) { documents := make([]ruleDocument, 0, len(paths)) for _, path := range paths { @@ -55,10 +59,13 @@ func loadRuleDocuments(ctx context.Context, paths []string, readFile func(string return documents, nil } +// discoverRuleFiles 执行默认规则文件发现策略。 func discoverRuleFiles(ctx context.Context, workdir string) ([]string, error) { return discoverRuleFilesWithFinder(ctx, workdir, findExactRuleFile) } +// discoverRuleFilesWithFinder 从当前目录向上遍历到根目录，收集匹配规则文件。 +// 返回顺序会反转为“从上层到下层”，让更顶层规则先出现。 func discoverRuleFilesWithFinder(ctx context.Context, workdir string, finder ruleFileFinder) ([]string, error) { workdir = strings.TrimSpace(workdir) if workdir == "" { @@ -98,6 +105,7 @@ func discoverRuleFilesWithFinder(ctx context.Context, workdir string, finder rul return paths, nil } +// findExactRuleFile 在单个目录内查找严格大小写匹配的 AGENTS.md。 func findExactRuleFile(dir string) (string, error) { entries, err := os.ReadDir(dir) if err != nil { @@ -119,6 +127,8 @@ func findExactRuleFile(dir string) (string, error) { return "", nil } +// renderProjectRulesSection 在总预算内渲染规则分节。 +// 当预算不足时，优先保留靠近工作目录的规则并标记截断提示。 func renderProjectRulesSection(documents []ruleDocument) promptSection { if len(documents) == 0 { return promptSection{} @@ -167,6 +177,7 @@ func renderProjectRulesSection(documents []ruleDocument) promptSection { } } +// renderRuleDocumentChunk 渲染单个规则文件块（不考虑总预算）。 func renderRuleDocumentChunk(document ruleDocument) string { var builder strings.Builder builder.WriteString("\n### ") @@ -184,6 +195,8 @@ func renderRuleDocumentChunk(document ruleDocument) string { return builder.String() } +// renderRuleDocumentChunkWithinBudget 在给定字符预算内渲染单个规则文件块。 +// 若预算不足以容纳标题，则直接返回空字符串。 func renderRuleDocumentChunkWithinBudget(document ruleDocument, budget int) string { if budget <= 0 { return "" @@ -224,6 +237,7 @@ func renderRuleDocumentChunkWithinBudget(document ruleDocument, budget int) stri return header + body.String() } +// truncateRunes 按 rune（而非字节）截断，避免 UTF-8 多字节字符被破坏。 func truncateRunes(input string, max int) (string, bool) { if max <= 0 { return "", input != "" @@ -236,6 +250,7 @@ func truncateRunes(input string, max int) (string, bool) { return string(runes[:max]), true } +// runeCount 统一使用 rune 数进行长度预算计算。 func runeCount(input string) int { return utf8.RuneCountInString(input) } diff --git a/internal/context/source_system.go b/internal/context/source_system.go index 6ee12e0d..67699b6d 100644 --- a/internal/context/source_system.go +++ b/internal/context/source_system.go @@ -8,8 +8,11 @@ import ( "strings" ) +// gitCommandRunner 抽象 git 命令执行，方便测试替身注入。 type gitCommandRunner func(ctx context.Context, workdir string, args ...string) (string, error) +// collectSystemState 聚合运行时环境信息，并尽力探测 git 状态。 +// 非上下文取消类 git 错误会被吞掉，以保证主流程可继续。 func collectSystemState(ctx context.Context, metadata Metadata, runner gitCommandRunner) (SystemState, error) { state := SystemState{ Workdir: strings.TrimSpace(metadata.Workdir), @@ -25,6 +28,7 @@ func collectSystemState(ctx context.Context, metadata Metadata, runner gitComman return state, nil } + // 当前分支读取失败时，不中断上下文构建；只有上下文取消才返回错误。 branch, err := runner(ctx, state.Workdir, "rev-parse", "--abbrev-ref", "HEAD") if err != nil { if isContextError(err) { @@ -32,6 +36,7 @@ func collectSystemState(ctx context.Context, metadata Metadata, runner gitComman } return state, nil } + // 通过 porcelain 输出是否为空判断工作区是否脏。 dirty, err := runner(ctx, state.Workdir, "status", "--porcelain") if err != nil { if isContextError(err) { @@ -48,6 +53,7 @@ func collectSystemState(ctx context.Context, metadata Metadata, runner gitComman return state, nil } +// renderSystemStateSection 将系统态转为结构化提示词文本。 func renderSystemStateSection(state SystemState) promptSection { lines := []string{ fmt.Sprintf("- workdir: `%s`", promptValue(state.Workdir)), @@ -72,6 +78,7 @@ func renderSystemStateSection(state SystemState) promptSection { } } +// runGitCommand 使用 `git -C ` 在指定目录执行 git 子命令。 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() @@ -81,6 +88,7 @@ func runGitCommand(ctx context.Context, workdir string, args ...string) (string, return string(output), nil } +// promptValue 把空值规范为 unknown，避免向提示词暴露空字符串。 func promptValue(value string) string { value = strings.TrimSpace(value) if value == "" { @@ -89,6 +97,7 @@ func promptValue(value string) string { return value } +// isContextError 识别可传播的上下文取消/超时错误。 func isContextError(err error) bool { return errors.Is(err, context.Canceled) || errors.Is(err, context.DeadlineExceeded) } diff --git a/internal/context/trim.go b/internal/context/trim.go index 8947e7f3..47fed25e 100644 --- a/internal/context/trim.go +++ b/internal/context/trim.go @@ -4,6 +4,9 @@ import "neo-code/internal/provider" const maxContextTurns = 10 +// trimMessages 将历史消息裁剪到最近 maxContextTurns 个“对话片段”。 +// 片段定义：普通单条消息算一个片段；assistant(tool_calls)+后续tool结果算一个整体片段。 +// 这样可以在压缩上下文时尽量保持 tool_call 与 tool_result 的配对完整性。 func trimMessages(messages []provider.Message) []provider.Message { if len(messages) <= maxContextTurns { return append([]provider.Message(nil), messages...) @@ -19,6 +22,7 @@ func trimMessages(messages []provider.Message) []provider.Message { start := i i++ + // assistant 触发工具后，连续的 tool 消息都归属到同一片段。 if messages[start].Role == provider.RoleAssistant && len(messages[start].ToolCalls) > 0 { for i < len(messages) && messages[i].Role == provider.RoleTool { i++ @@ -32,6 +36,7 @@ func trimMessages(messages []provider.Message) []provider.Message { return append([]provider.Message(nil), messages...) } + // 仅保留最近 N 个片段，从片段起点切割，避免破坏消息对齐关系。 start := spans[len(spans)-maxContextTurns].start return append([]provider.Message(nil), messages[start:]...) } diff --git a/internal/context/types.go b/internal/context/types.go index 6eb516dd..b7fbc42f 100644 --- a/internal/context/types.go +++ b/internal/context/types.go @@ -7,17 +7,20 @@ import ( ) // Builder builds the provider-facing context for a single model round. +// Builder 抽象“每轮模型调用前”的上下文构建过程。 type Builder interface { Build(ctx context.Context, input BuildInput) (BuildResult, error) } // BuildInput contains the runtime state needed to assemble model context. +// BuildInput 是构建上下文所需输入，包含消息历史与运行时元信息。 type BuildInput struct { Messages []provider.Message Metadata Metadata } // BuildResult is the provider-facing context produced for a single round. +// BuildResult 是可直接喂给 provider 的上下文结果。 type BuildResult struct { SystemPrompt string Messages []provider.Message