From 01f020a7249ba5d30f5f1f6df23d976a5e1dfa4e Mon Sep 17 00:00:00 2001 From: phantom5099 <1011668688@qq.com> Date: Mon, 13 Apr 2026 00:35:14 +0800 Subject: [PATCH 1/5] =?UTF-8?q?refactor(runtime):=E4=BC=98=E5=8C=96runtime?= =?UTF-8?q?=E5=86=85=E9=83=A8=E8=81=8C=E8=B4=A3=E4=B8=8D=E6=B8=85=E9=97=AE?= =?UTF-8?q?=E9=A2=98?= MIME-Version: 1.0 Content-Type: text/plain; charset=UTF-8 Content-Transfer-Encoding: 8bit --- AGENTS.md | 2 +- docs/context-compact.md | 4 +- docs/guides/configuration.md | 2 +- docs/runtime-provider-event-flow.md | 29 +- internal/context/builder.go | 10 +- internal/context/builder_test.go | 24 +- internal/context/source_rules.go | 19 + internal/context/source_rules_test.go | 13 +- internal/context/types.go | 6 +- internal/provider/streaming/accumulator.go | 115 +++ internal/provider/streaming/handler.go | 64 ++ internal/runtime/approval/broker.go | 83 ++ internal/runtime/approval/broker_test.go | 95 +++ internal/runtime/approval/decision.go | 10 + internal/runtime/compact.go | 91 +-- internal/runtime/compact_generator.go | 52 +- internal/runtime/errors.go | 48 ++ internal/runtime/permission.go | 266 +++---- internal/runtime/permission_test.go | 117 +-- internal/runtime/provider_stream.go | 89 +++ internal/runtime/run.go | 272 +++++++ internal/runtime/runtime.go | 709 ++---------------- internal/runtime/runtime_gap_coverage_test.go | 288 +++++++ .../runtime/runtime_internal_helpers_test.go | 241 ++++++ .../runtime_remaining_branches_test.go | 478 ++++++++++++ internal/runtime/runtime_test.go | 215 ++---- internal/runtime/session_mutation.go | 67 ++ internal/runtime/state.go | 84 +++ internal/runtime/toolexec.go | 61 ++ internal/runtime/workdir_branch_test.go | 4 +- internal/tui/core/app/permission_prompt.go | 13 +- .../tui/core/app/permission_prompt_test.go | 15 +- internal/tui/core/app/update.go | 5 +- .../tui/core/app/update_permission_test.go | 25 +- internal/tui/core/app/update_test.go | 21 +- internal/tui/services/services_test.go | 7 +- 36 files changed, 2446 insertions(+), 1198 deletions(-) create mode 100644 internal/provider/streaming/accumulator.go create mode 100644 internal/provider/streaming/handler.go create mode 100644 internal/runtime/approval/broker.go create mode 100644 internal/runtime/approval/broker_test.go create mode 100644 internal/runtime/approval/decision.go create mode 100644 internal/runtime/errors.go create mode 100644 internal/runtime/provider_stream.go create mode 100644 internal/runtime/run.go create mode 100644 internal/runtime/runtime_gap_coverage_test.go create mode 100644 internal/runtime/runtime_internal_helpers_test.go create mode 100644 internal/runtime/runtime_remaining_branches_test.go create mode 100644 internal/runtime/session_mutation.go create mode 100644 internal/runtime/state.go create mode 100644 internal/runtime/toolexec.go diff --git a/AGENTS.md b/AGENTS.md index 79579751..b0c7a887 100644 --- a/AGENTS.md +++ b/AGENTS.md @@ -66,7 +66,7 @@ - token 累积记录与事件发射 - 自动压缩触发与重置逻辑 - 修改 `context` 时，重点覆盖： - - Build 输入输出契约（含 Metadata 新字段、ShouldAutoCompact 决策） + - Build 输入输出契约（含 Metadata 新字段、AutoCompactSuggested 决策） - micro compact 策略 - 消息裁剪边界 - AGENTS.md 加载与截断 diff --git a/docs/context-compact.md b/docs/context-compact.md index f2a60482..ade2e732 100644 --- a/docs/context-compact.md +++ b/docs/context-compact.md @@ -5,7 +5,7 @@ ## 概览 - runtime 已接入手动 compact、基于 token 阈值的自动 compact，以及 provider 上下文过长后的 `reactive` compact 自动恢复。 -- `internal/context/compact` 支持 `manual` 与 `reactive` 两种 mode。 +- `internal/context/compact` 支持 `manual`、`auto` 与 `reactive` 三种 mode。 - 用户通过 `/compact` 对当前会话执行一次上下文压缩。 - compact 前会先写入完整 transcript，随后生成并校验 compact summary，再回写会话消息。 @@ -40,7 +40,7 @@ context: ## 自动压缩 -当 `auto_compact.enabled` 为 `true` 时，runtime 在每次调用 `context.Builder.Build()` 时将当前 token 累计值传入 Metadata，context 模块通过比较累计值与阈值在 `BuildResult.ShouldAutoCompact` 中返回压缩建议。runtime 读取建议后调用现有 compact 管线执行压缩；token 计数的重置与持久化语义统一见 [Session 持久化设计](./session-persistence-design.md)。 +当 `auto_compact.enabled` 为 `true` 时，runtime 在每次调用 `context.Builder.Build()` 时将当前 token 累计值传入 Metadata，context 模块通过比较累计值与阈值在 `BuildResult.AutoCompactSuggested` 中返回压缩建议。runtime 读取建议后调用现有 compact 管线执行压缩；token 计数的重置与持久化语义统一见 [Session 持久化设计](./session-persistence-design.md)。 设计原则： - **context 拥有压缩决策权**，runtime 只做编排执行。 diff --git a/docs/guides/configuration.md b/docs/guides/configuration.md index fcc7293a..f7beab33 100644 --- a/docs/guides/configuration.md +++ b/docs/guides/configuration.md @@ -171,7 +171,7 @@ $env:GEMINI_API_KEY = "AI..." `config.yaml` 里的 `selected_provider/current_model` 表达的是“用户上次保存的选择状态”。 -启动时系统还会基于当前 provider、driver 支持情况和模型目录快照执行选择修正。因此需要区分两件事： +启动时系统还会进行选择校验与必要修正；若 driver 不受支持会报错并中止。因此需要区分两件事： - 配置快照结构合法 - 当前选择已经可直接运行 diff --git a/docs/runtime-provider-event-flow.md b/docs/runtime-provider-event-flow.md index 8b4cde31..8c4eb523 100644 --- a/docs/runtime-provider-event-flow.md +++ b/docs/runtime-provider-event-flow.md @@ -4,11 +4,18 @@ 当前 runtime 对外暴露一组小而稳定的事件： +- `user_message` - `agent_chunk` - `agent_done` +- `tool_call_thinking` - `tool_start` +- `tool_chunk` - `tool_result` +- `run_canceled` - `error` +- `provider_retry` +- `permission_request` +- `permission_resolved` - `token_usage` - `compact_start` - `compact_done` @@ -19,15 +26,20 @@ 1. 加载目标会话或创建新会话。 2. 追加最新的用户消息。 3. 读取最新配置快照。 -4. 解析当前 provider 配置并构建 provider 实例。 -5. 调用 `context.Builder` 生成本轮请求使用的 `system prompt` 和消息上下文。 -6. 如命中 token 阈值自动压缩建议，则先执行一次 compact，再继续构造请求。 +4. 调用 `context.Builder` 生成本轮请求使用的 `system prompt` 和消息上下文。 +5. 如命中 token 阈值自动压缩建议，则先执行一次 compact，再在同一轮内重建请求。 +6. 冻结当前 turn 的 `provider / model / tools / workdir / request` 快照。 7. 调用 `Provider.Generate`，并把流式事件桥接给 TUI。 -8. 如 provider 返回“上下文过长”错误，则触发一次 `reactive` compact，并仅重试一次当前请求。 +8. 如 provider 返回“上下文过长”错误，则触发一次 `reactive` compact，并仅在同一 turn 内重建一次当前请求。 9. 保存 assistant 完整回复。 10. 执行返回的工具调用，并保存每一个工具结果。 11. 如果仍需继续推理，则进入下一轮；否则结束。 +补充约束： +- 同一 turn 内的 provider retry 只重放冻结后的 turn 快照，不会重新读取配置。 +- `auto compact` 与 `reactive compact` 都不额外消耗 reasoning turn。 +- 权限审批等待由 `internal/runtime/approval` 负责 request 生命周期，runtime 自己负责事件发射与 tool 重试编排。 + ### Context Builder 输入与职责 - `runtime` 只向 `context.Builder` 传递本轮所需元数据： @@ -44,7 +56,7 @@ - 从 `workdir` 向上发现的 `AGENTS.md` - 系统状态摘要（`workdir` / `shell` / `provider` / `model` / git branch / git dirty） - 裁剪后的历史消息 - - 自动压缩决策（`BuildResult.ShouldAutoCompact`） + - 自动压缩决策（`BuildResult.AutoCompactSuggested`） - `runtime` 不直接读取规则文件，也不直接查询 git 状态。 - `provider` 只消费最终生成的 `SystemPrompt`、消息列表和工具 schema，不感知上下文来源。 @@ -66,9 +78,14 @@ ## 流式桥接 - Provider 发出 `StreamEvent` -- runtime 将其转换成 `RuntimeEvent` +- `internal/provider/streaming` 统一累积文本、tool call 增量和 `message_done` +- runtime 将累积过程映射成 `RuntimeEvent` - TUI 使用 Bubble Tea `Cmd` 监听事件，并在处理完成后继续订阅 +同一套流式累积逻辑同时复用于： +- 普通 `Run()` 的 assistant 回复收敛 +- compact summary 生成阶段的 provider 输出消费 + ## Token 计量 runtime 在转发 provider 流式事件时，从 `MessageDone` 事件中提取 `Usage`（`InputTokens`、`OutputTokens`），累积到会话级计数器，并发出 `token_usage` 事件供 TUI 消费。 diff --git a/internal/context/builder.go b/internal/context/builder.go index de8ea511..acf2f29e 100644 --- a/internal/context/builder.go +++ b/internal/context/builder.go @@ -47,8 +47,8 @@ func NewBuilderWithMemo(policies MicroCompactPolicySource, memoSource SectionSou sources = append(sources, memoSource) } return &DefaultBuilder{ - promptSources: sources, - trimPolicy: spanMessageTrimPolicy{}, + promptSources: sources, + trimPolicy: spanMessageTrimPolicy{}, microCompactPolicies: policies, } } @@ -77,9 +77,9 @@ func (b *DefaultBuilder) Build(ctx context.Context, input BuildInput) (BuildResu input.Metadata.SessionInputTokens >= input.Compact.AutoCompactThreshold return BuildResult{ - SystemPrompt: composeSystemPrompt(sections...), - Messages: applyReadTimeContextProjection(trimPolicy.Trim(input.Messages), input.Compact, b.microCompactPolicies), - ShouldAutoCompact: shouldAutoCompact, + SystemPrompt: composeSystemPrompt(sections...), + Messages: applyReadTimeContextProjection(trimPolicy.Trim(input.Messages), input.Compact, b.microCompactPolicies), + AutoCompactSuggested: shouldAutoCompact, }, nil } diff --git a/internal/context/builder_test.go b/internal/context/builder_test.go index 700ef159..c21a2f3a 100644 --- a/internal/context/builder_test.go +++ b/internal/context/builder_test.go @@ -551,7 +551,7 @@ func TestTrimMessagesBoundaries(t *testing.T) { } } -func TestBuildShouldAutoCompactDisabled(t *testing.T) { +func TestBuildAutoCompactSuggestedDisabled(t *testing.T) { t.Parallel() builder := NewBuilder() @@ -566,12 +566,12 @@ func TestBuildShouldAutoCompactDisabled(t *testing.T) { if err != nil { t.Fatalf("Build() error = %v", err) } - if result.ShouldAutoCompact { - t.Fatalf("expected ShouldAutoCompact false when threshold is 0") + if result.AutoCompactSuggested { + t.Fatalf("expected AutoCompactSuggested false when threshold is 0") } } -func TestBuildShouldAutoCompactBelowThreshold(t *testing.T) { +func TestBuildAutoCompactSuggestedBelowThreshold(t *testing.T) { t.Parallel() builder := NewBuilder() @@ -586,12 +586,12 @@ func TestBuildShouldAutoCompactBelowThreshold(t *testing.T) { if err != nil { t.Fatalf("Build() error = %v", err) } - if result.ShouldAutoCompact { - t.Fatalf("expected ShouldAutoCompact false when tokens below threshold") + if result.AutoCompactSuggested { + t.Fatalf("expected AutoCompactSuggested false when tokens below threshold") } } -func TestBuildShouldAutoCompactAtThreshold(t *testing.T) { +func TestBuildAutoCompactSuggestedAtThreshold(t *testing.T) { t.Parallel() builder := NewBuilder() @@ -606,12 +606,12 @@ func TestBuildShouldAutoCompactAtThreshold(t *testing.T) { if err != nil { t.Fatalf("Build() error = %v", err) } - if !result.ShouldAutoCompact { - t.Fatalf("expected ShouldAutoCompact true when tokens equal threshold") + if !result.AutoCompactSuggested { + t.Fatalf("expected AutoCompactSuggested true when tokens equal threshold") } } -func TestBuildShouldAutoCompactAboveThreshold(t *testing.T) { +func TestBuildAutoCompactSuggestedAboveThreshold(t *testing.T) { t.Parallel() builder := NewBuilder() @@ -626,8 +626,8 @@ func TestBuildShouldAutoCompactAboveThreshold(t *testing.T) { if err != nil { t.Fatalf("Build() error = %v", err) } - if !result.ShouldAutoCompact { - t.Fatalf("expected ShouldAutoCompact true when tokens above threshold") + if !result.AutoCompactSuggested { + t.Fatalf("expected AutoCompactSuggested true when tokens above threshold") } } diff --git a/internal/context/source_rules.go b/internal/context/source_rules.go index 83323147..94ab90d3 100644 --- a/internal/context/source_rules.go +++ b/internal/context/source_rules.go @@ -2,7 +2,9 @@ package context import ( "context" + "errors" "fmt" + "io/fs" "os" "path/filepath" "strings" @@ -216,6 +218,9 @@ func discoverRuleFilesWithFinder(ctx context.Context, workdir string, finder rul match, err := finder(dir) if err != nil { + if isRuleDiscoveryPermissionError(err) { + break + } return nil, fmt.Errorf("context: discover rule file in %s: %w", dir, err) } if match != "" { @@ -236,6 +241,20 @@ func discoverRuleFilesWithFinder(ctx context.Context, workdir string, finder rul return paths, nil } +// isRuleDiscoveryPermissionError 判断规则发现失败是否由权限限制导致。 +// 在沙箱或受限目录场景下，遇到无权限读取的父目录时应停止继续向上探测，而不是让整个上下文构建失败。 +func isRuleDiscoveryPermissionError(err error) bool { + if err == nil { + return false + } + if os.IsPermission(err) || errors.Is(err, fs.ErrPermission) || errors.Is(err, os.ErrPermission) { + return true + } + + lower := strings.ToLower(err.Error()) + return strings.Contains(lower, "permission denied") || strings.Contains(lower, "access is denied") +} + // findExactRuleFile 只匹配大小写完全一致的 AGENTS.md，避免误读同名变体。 func findExactRuleFile(dir string) (string, error) { entries, err := os.ReadDir(dir) diff --git a/internal/context/source_rules_test.go b/internal/context/source_rules_test.go index 9b90aca7..687cae9d 100644 --- a/internal/context/source_rules_test.go +++ b/internal/context/source_rules_test.go @@ -3,6 +3,7 @@ package context import ( "context" "errors" + "fmt" "os" "path/filepath" "strings" @@ -115,7 +116,7 @@ func TestLoadRuleDocumentsReturnsReadError(t *testing.T) { } } -func TestDiscoverRuleFilesReturnsDirectoryReadError(t *testing.T) { +func TestDiscoverRuleFilesStopsTraversalOnPermissionDenied(t *testing.T) { t.Parallel() root := t.TempDir() @@ -133,7 +134,7 @@ func TestDiscoverRuleFilesReturnsDirectoryReadError(t *testing.T) { t.Fatalf("write local rules: %v", err) } - permissionErr := errors.New("permission denied") + permissionErr := fmt.Errorf("wrapped permission: %w", os.ErrPermission) paths, err := discoverRuleFilesWithFinder(context.Background(), nested, func(dir string) (string, error) { switch dir { case nested: @@ -146,11 +147,11 @@ func TestDiscoverRuleFilesReturnsDirectoryReadError(t *testing.T) { return "", nil } }) - if err == nil || !strings.Contains(err.Error(), permissionErr.Error()) { - t.Fatalf("expected discoverRuleFilesWithFinder() to return permission error, got %v", err) + if err != nil { + t.Fatalf("discoverRuleFilesWithFinder() error = %v", err) } - if paths != nil { - t.Fatalf("expected no paths on discovery failure, got %+v", paths) + if len(paths) != 1 || paths[0] != localRules { + t.Fatalf("expected discovery to stop after permission denial, got %+v", paths) } } diff --git a/internal/context/types.go b/internal/context/types.go index 6fadc524..620d4ba1 100644 --- a/internal/context/types.go +++ b/internal/context/types.go @@ -21,9 +21,9 @@ type BuildInput struct { // BuildResult is the provider-facing context produced for a single round. type BuildResult struct { - SystemPrompt string - Messages []providertypes.Message - ShouldAutoCompact bool + SystemPrompt string + Messages []providertypes.Message + AutoCompactSuggested bool } // MicroCompactPolicySource 定义 context 读取工具 micro compact 策略的最小依赖。 diff --git a/internal/provider/streaming/accumulator.go b/internal/provider/streaming/accumulator.go new file mode 100644 index 00000000..8ec35b26 --- /dev/null +++ b/internal/provider/streaming/accumulator.go @@ -0,0 +1,115 @@ +package streaming + +import ( + "fmt" + "sort" + "strings" + + providertypes "neo-code/internal/provider/types" +) + +// Accumulator 负责按 provider 流式事件累积 assistant 文本与工具调用。 +type Accumulator struct { + content strings.Builder + toolCalls map[int]*providertypes.ToolCall + messageDone bool +} + +// NewAccumulator 创建一个空的流式消息累积器。 +func NewAccumulator() *Accumulator { + return &Accumulator{ + toolCalls: make(map[int]*providertypes.ToolCall), + } +} + +// MessageDone 返回当前流是否已经收到 message_done 事件。 +func (a *Accumulator) MessageDone() bool { + return a != nil && a.messageDone +} + +// BuildMessage 将当前累积状态收敛为最终 assistant 消息。 +func (a *Accumulator) BuildMessage() (providertypes.Message, error) { + if a == nil { + return providertypes.Message{Role: providertypes.RoleAssistant}, nil + } + + ordered := make([]int, 0, len(a.toolCalls)) + for index := range a.toolCalls { + ordered = append(ordered, index) + } + sort.Ints(ordered) + + message := providertypes.Message{ + Role: providertypes.RoleAssistant, + Content: a.content.String(), + } + for _, index := range ordered { + call := a.toolCalls[index] + if call == nil { + continue + } + if strings.TrimSpace(call.ID) == "" { + return providertypes.Message{}, fmt.Errorf("runtime: provider emitted tool call %d without id", index) + } + if strings.TrimSpace(call.Name) == "" { + return providertypes.Message{}, fmt.Errorf("runtime: provider emitted tool call %d without name", index) + } + message.ToolCalls = append(message.ToolCalls, *call) + } + return message, nil +} + +// AccumulateText 负责累积文本增量。 +func (a *Accumulator) AccumulateText(text string) { + if a == nil { + return + } + a.content.WriteString(text) +} + +// ensureToolCall 返回指定索引的工具调用占位对象。 +func (a *Accumulator) ensureToolCall(index int) *providertypes.ToolCall { + if a == nil { + return nil + } + call, exists := a.toolCalls[index] + if !exists { + call = &providertypes.ToolCall{} + a.toolCalls[index] = call + } + return call +} + +// AccumulateToolCallStart 负责合并工具调用的起始元数据。 +func (a *Accumulator) AccumulateToolCallStart(index int, id string, name string) { + call := a.ensureToolCall(index) + if call == nil { + return + } + if strings.TrimSpace(id) != "" { + call.ID = id + } + if strings.TrimSpace(name) != "" { + call.Name = name + } +} + +// AccumulateToolCallDelta 负责累积工具调用参数增量。 +func (a *Accumulator) AccumulateToolCallDelta(index int, id string, argumentsDelta string) { + call := a.ensureToolCall(index) + if call == nil { + return + } + if strings.TrimSpace(id) != "" { + call.ID = id + } + call.Arguments += argumentsDelta +} + +// MarkMessageDone 标记当前流已经收到最终完成事件。 +func (a *Accumulator) MarkMessageDone() { + if a == nil { + return + } + a.messageDone = true +} diff --git a/internal/provider/streaming/handler.go b/internal/provider/streaming/handler.go new file mode 100644 index 00000000..9280b98f --- /dev/null +++ b/internal/provider/streaming/handler.go @@ -0,0 +1,64 @@ +package streaming + +import ( + "fmt" + + providertypes "neo-code/internal/provider/types" +) + +// Hooks 描述 runtime 在消费 provider 流时可选的回调挂点。 +type Hooks struct { + OnTextDelta func(string) + OnToolCallStart func(providertypes.ToolCallStartPayload) + OnMessageDone func(providertypes.MessageDonePayload) +} + +// HandleEvent 解析并应用单个 provider 流式事件。 +func HandleEvent(event providertypes.StreamEvent, acc *Accumulator, hooks Hooks) error { + switch event.Type { + case providertypes.StreamEventTextDelta: + payload, err := event.TextDeltaValue() + if err != nil { + return err + } + if hooks.OnTextDelta != nil { + hooks.OnTextDelta(payload.Text) + } + if acc != nil { + acc.AccumulateText(payload.Text) + } + case providertypes.StreamEventToolCallStart: + payload, err := event.ToolCallStartValue() + if err != nil { + return err + } + if hooks.OnToolCallStart != nil { + hooks.OnToolCallStart(payload) + } + if acc != nil { + acc.AccumulateToolCallStart(payload.Index, payload.ID, payload.Name) + } + case providertypes.StreamEventToolCallDelta: + payload, err := event.ToolCallDeltaValue() + if err != nil { + return err + } + if acc != nil { + acc.AccumulateToolCallDelta(payload.Index, payload.ID, payload.ArgumentsDelta) + } + case providertypes.StreamEventMessageDone: + payload, err := event.MessageDoneValue() + if err != nil { + return err + } + if acc != nil { + acc.MarkMessageDone() + } + if hooks.OnMessageDone != nil { + hooks.OnMessageDone(payload) + } + default: + return fmt.Errorf("runtime: unsupported provider stream event type %q", event.Type) + } + return nil +} diff --git a/internal/runtime/approval/broker.go b/internal/runtime/approval/broker.go new file mode 100644 index 00000000..cdaeacdb --- /dev/null +++ b/internal/runtime/approval/broker.go @@ -0,0 +1,83 @@ +package approval + +import ( + "errors" + "fmt" + "sync" +) + +type pendingRequest struct { + resultCh chan Decision + submitted bool +} + +// Broker 负责管理运行期间待审批请求的生命周期。 +type Broker struct { + mu sync.Mutex + nextID uint64 + pending map[string]*pendingRequest +} + +// NewBroker 创建一个空的审批请求 broker。 +func NewBroker() *Broker { + return &Broker{ + pending: make(map[string]*pendingRequest), + } +} + +// Open 注册一个新的待审批请求，并返回 request id 与结果通道。 +func (b *Broker) Open() (string, chan Decision, error) { + if b == nil { + return "", nil, errors.New("runtime: approval broker is nil") + } + + b.mu.Lock() + defer b.mu.Unlock() + + b.nextID++ + requestID := fmt.Sprintf("perm-%d", b.nextID) + request := &pendingRequest{ + resultCh: make(chan Decision, 1), + } + b.pending[requestID] = request + return requestID, request.resultCh, nil +} + +// Resolve 向指定请求提交审批结果；重复提交会被安全忽略。 +func (b *Broker) Resolve(requestID string, decision Decision) error { + if b == nil { + return errors.New("runtime: approval broker is nil") + } + + b.mu.Lock() + request := b.pending[requestID] + if request == nil { + b.mu.Unlock() + return fmt.Errorf("runtime: permission request %q not found", requestID) + } + if request.submitted { + b.mu.Unlock() + return nil + } + request.submitted = true + resultCh := request.resultCh + b.mu.Unlock() + + select { + case resultCh <- decision: + return nil + default: + return nil + } +} + +// Close 清理指定请求，避免后续误用过期 request id。 +func (b *Broker) Close(requestID string) { + if b == nil { + return + } + + b.mu.Lock() + defer b.mu.Unlock() + delete(b.pending, requestID) +} diff --git a/internal/runtime/approval/broker_test.go b/internal/runtime/approval/broker_test.go new file mode 100644 index 00000000..2157eba2 --- /dev/null +++ b/internal/runtime/approval/broker_test.go @@ -0,0 +1,95 @@ +package approval + +import ( + "testing" +) + +func TestBrokerNilReceiverBranches(t *testing.T) { + t.Parallel() + + var broker *Broker + if _, _, err := broker.Open(); err == nil { + t.Fatalf("expected open on nil broker to fail") + } + if err := broker.Resolve("perm-1", DecisionAllowOnce); err == nil { + t.Fatalf("expected resolve on nil broker to fail") + } + broker.Close("perm-1") // should be no-op +} + +func TestBrokerOpenResolveCloseFlow(t *testing.T) { + t.Parallel() + + broker := NewBroker() + requestID, resultCh, err := broker.Open() + if err != nil { + t.Fatalf("Open() error = %v", err) + } + if requestID == "" { + t.Fatalf("expected non-empty request id") + } + + if err := broker.Resolve(requestID, DecisionAllowSession); err != nil { + t.Fatalf("Resolve() error = %v", err) + } + + select { + case decision := <-resultCh: + if decision != DecisionAllowSession { + t.Fatalf("expected %q, got %q", DecisionAllowSession, decision) + } + default: + t.Fatalf("expected decision to be delivered") + } + + // duplicate submit should be idempotent and non-blocking + if err := broker.Resolve(requestID, DecisionReject); err != nil { + t.Fatalf("duplicate Resolve() error = %v", err) + } + + broker.Close(requestID) + if err := broker.Resolve(requestID, DecisionAllowOnce); err == nil { + t.Fatalf("expected resolve after close to fail") + } +} + +func TestBrokerResolveUnknownRequest(t *testing.T) { + t.Parallel() + + broker := NewBroker() + if err := broker.Resolve("perm-missing", DecisionAllowOnce); err == nil { + t.Fatalf("expected missing request error") + } +} + +func TestBrokerResolveChannelFullStillReturnsNil(t *testing.T) { + t.Parallel() + + broker := NewBroker() + requestID, resultCh, err := broker.Open() + if err != nil { + t.Fatalf("Open() error = %v", err) + } + + // first resolve fills the buffered channel (capacity 1) + if err := broker.Resolve(requestID, DecisionAllowOnce); err != nil { + t.Fatalf("first Resolve() error = %v", err) + } + + // simulate an already-submitted request whose channel is full, + // to cover select default branch in Resolve. + broker.mu.Lock() + broker.pending[requestID] = &pendingRequest{resultCh: make(chan Decision, 1)} + broker.pending[requestID].resultCh <- DecisionAllowOnce + broker.mu.Unlock() + + if err := broker.Resolve(requestID, DecisionAllowSession); err != nil { + t.Fatalf("Resolve() on full channel should still return nil, got %v", err) + } + + // ensure original channel remains usable for read and not leaked behaviorally + select { + case <-resultCh: + default: + } +} diff --git a/internal/runtime/approval/decision.go b/internal/runtime/approval/decision.go new file mode 100644 index 00000000..0bd014fb --- /dev/null +++ b/internal/runtime/approval/decision.go @@ -0,0 +1,10 @@ +package approval + +// Decision 表示审批请求的最终处理结果。 +type Decision string + +const ( + DecisionAllowOnce Decision = "allow_once" + DecisionAllowSession Decision = "allow_session" + DecisionReject Decision = "reject" +) diff --git a/internal/runtime/compact.go b/internal/runtime/compact.go index 210af144..b7a1227d 100644 --- a/internal/runtime/compact.go +++ b/internal/runtime/compact.go @@ -12,13 +12,13 @@ import ( agentsession "neo-code/internal/session" ) -// CompactInput 描述一次手动 compact 请求所需的最小输入。 +// CompactInput 描述一次手动 compact 请求的最小输入。 type CompactInput struct { SessionID string RunID string } -// CompactResult 汇总 compact 执行后的外部可见结果。 +// CompactResult 汇总一次 compact 完成后的外部可见结果，统一作为返回值和事件 payload 使用。 type CompactResult struct { Applied bool BeforeChars int @@ -29,24 +29,35 @@ type CompactResult struct { TranscriptPath string } -// CompactDonePayload 是 compact 完成事件向上层暴露的摘要信息。 -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"` +// fromCompactResult 将 contextcompact.Result 映射为 runtime.CompactResult。 +func fromCompactResult(result contextcompact.Result) CompactResult { + 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, + } } -// CompactErrorPayload 是 compact 失败事件向上层暴露的错误信息。 +// CompactErrorPayload 是 compact_error 事件对外暴露的错误信息。 type CompactErrorPayload struct { TriggerMode string `json:"trigger_mode"` Message string `json:"message"` } -// Compact 串行执行手动 compact，并返回本次压缩的统计结果。 +// compactErrorPolicy 描述 compact 失败后是立即中断还是仅记录事件后继续运行。 +type compactErrorPolicy uint8 + +const ( + compactErrorStrict compactErrorPolicy = iota + compactErrorBestEffort +) + +// Compact 串行执行一次手动 compact，并返回本次压缩统计信息。 +// 会话级锁确保同一会话的 Run 和 Compact 互斥，不同会话可并行。 func (s *Service) Compact(ctx context.Context, input CompactInput) (CompactResult, error) { if err := ctx.Err(); err != nil { return CompactResult{}, err @@ -55,8 +66,9 @@ func (s *Service) Compact(ctx context.Context, input CompactInput) (CompactResul return CompactResult{}, errors.New("runtime: compact session_id is empty") } - s.operationMu.Lock() - defer s.operationMu.Unlock() + sessionMu := s.acquireSessionLock(input.SessionID) + sessionMu.Lock() + defer sessionMu.Unlock() cfg := s.configManager.Get() session, err := s.sessionStore.Load(ctx, input.SessionID) @@ -64,20 +76,12 @@ func (s *Service) Compact(ctx context.Context, input CompactInput) (CompactResul return CompactResult{}, err } - session, result, err := s.runCompactForSession(ctx, input.RunID, session, cfg, contextcompact.ModeManual, true) + session, result, err := s.runCompactForSession(ctx, input.RunID, session, cfg, contextcompact.ModeManual, compactErrorStrict) 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 + return fromCompactResult(result), nil } // runCompactForSession 负责发出 compact 事件、调用 runner，并在成功后回写会话。 @@ -87,7 +91,7 @@ func (s *Service) runCompactForSession( session agentsession.Session, cfg config.Config, mode contextcompact.Mode, - failOnError bool, + errorPolicy compactErrorPolicy, ) (agentsession.Session, contextcompact.Result, error) { runner := s.compactRunner if runner == nil { @@ -98,7 +102,7 @@ func (s *Service) runCompactForSession( TriggerMode: string(mode), Message: err.Error(), }) - if failOnError { + if errorPolicy == compactErrorStrict { return session, contextcompact.Result{}, err } return session, contextcompact.Result{}, nil @@ -111,7 +115,7 @@ func (s *Service) runCompactForSession( result, err := runner.Run(ctx, contextcompact.Input{ Mode: mode, SessionID: session.ID, - Workdir: effectiveSessionWorkdir(session.Workdir, cfg.Workdir), + Workdir: agentsession.EffectiveWorkdir(session.Workdir, cfg.Workdir), Messages: session.Messages, Config: cfg.Context.Compact, }) @@ -120,7 +124,7 @@ func (s *Service) runCompactForSession( TriggerMode: string(mode), Message: err.Error(), }) - if failOnError { + if errorPolicy == compactErrorStrict { return session, contextcompact.Result{}, err } return session, contextcompact.Result{}, nil @@ -137,39 +141,18 @@ func (s *Service) runCompactForSession( Message: err.Error(), }) session.Messages = originalMessages - if failOnError { + if errorPolicy == compactErrorStrict { 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: string(mode), - TranscriptID: result.TranscriptID, - TranscriptPath: result.TranscriptPath, - } - s.emit(ctx, EventCompactDone, runID, session.ID, donePayload) - + s.emit(ctx, EventCompactDone, runID, session.ID, fromCompactResult(result)) return session, result, nil } -// resetSessionTokenTotals 在 compact 成功后同步清零内存计数器与会话累计 token。 -func (s *Service) resetSessionTokenTotals(session *agentsession.Session) { - s.sessionInputTokens = 0 - s.sessionOutputTokens = 0 - if session == nil { - return - } - session.TokenInputTotal = 0 - session.TokenOutputTotal = 0 -} - -// defaultCompactRunner 为手动 compact 选择摘要生成器并构造默认 runner。 +// defaultCompactRunner 为 runtime 懒加载默认 compact runner。 func (s *Service) defaultCompactRunner(session agentsession.Session, cfg config.Config) (contextcompact.Runner, error) { resolvedProvider, model, err := resolveCompactProviderSelection(session, cfg) if err != nil { @@ -178,7 +161,7 @@ func (s *Service) defaultCompactRunner(session agentsession.Session, cfg config. return contextcompact.NewRunner(newCompactSummaryGenerator(s.providerFactory, resolvedProvider.ToRuntimeConfig(), model)), nil } -// resolveCompactProviderSelection 优先复用会话记录的 provider/model，缺失时回退到当前配置。 +// resolveCompactProviderSelection 优先复用会话最近成功运行时记录的 provider 与 model。 func resolveCompactProviderSelection(session agentsession.Session, cfg config.Config) (config.ResolvedProviderConfig, string, error) { sessionProvider := strings.TrimSpace(session.Provider) sessionModel := strings.TrimSpace(session.Model) diff --git a/internal/runtime/compact_generator.go b/internal/runtime/compact_generator.go index 061e1f6b..9b76818e 100644 --- a/internal/runtime/compact_generator.go +++ b/internal/runtime/compact_generator.go @@ -3,12 +3,12 @@ package runtime import ( "context" "errors" - "fmt" "strings" agentcontext "neo-code/internal/context" contextcompact "neo-code/internal/context/compact" "neo-code/internal/provider" + "neo-code/internal/provider/streaming" providertypes "neo-code/internal/provider/types" ) @@ -30,6 +30,7 @@ func newCompactSummaryGenerator( } } +// Generate 使用冻结后的 provider 配置为 compact 生成语义摘要。 func (g *compactSummaryGenerator) Generate(ctx context.Context, input contextcompact.SummaryInput) (string, error) { if err := ctx.Err(); err != nil { return "", err @@ -58,58 +59,19 @@ func (g *compactSummaryGenerator) Generate(ctx context.Context, input contextcom return "", err } - // 使用流式事件通道收集 compact 摘要响应。 - streamEvents := make(chan providertypes.StreamEvent, 32) - streamDone := make(chan error, 1) - acc := newStreamAccumulator() - - go func() { - var streamErr error - defer func() { - streamDone <- streamErr - }() - - for { - select { - case event, ok := <-streamEvents: - if !ok { - return - } - if err := handleProviderStreamEvent(event, acc, nil, nil, nil); err != nil && streamErr == nil { - // 记录首个协议错误后继续排空事件通道，避免 provider 在后续发送时阻塞。 - streamErr = err - } - case <-ctx.Done(): - return - } - } - }() - - err = modelProvider.Generate(ctx, providertypes.GenerateRequest{ + outcome := generateStreamingMessage(ctx, modelProvider, providertypes.GenerateRequest{ Model: g.model, SystemPrompt: prompt.SystemPrompt, Messages: []providertypes.Message{{ Role: providertypes.RoleUser, Content: prompt.UserPrompt, }}, - }, streamEvents) - close(streamEvents) - streamErr := <-streamDone - - if err != nil { - return "", err - } - if streamErr != nil { - return "", streamErr - } - if !acc.messageDone { - return "", fmt.Errorf("%w: provider stream ended without message_done event", provider.ErrStreamInterrupted) + }, streaming.Hooks{}) + if outcome.err != nil { + return "", outcome.err } - message, err := acc.buildMessage() - if err != nil { - return "", err - } + message := outcome.message if len(message.ToolCalls) > 0 { return "", errors.New("runtime: compact summary response must not contain tool calls") } diff --git a/internal/runtime/errors.go b/internal/runtime/errors.go new file mode 100644 index 00000000..c0ded3b7 --- /dev/null +++ b/internal/runtime/errors.go @@ -0,0 +1,48 @@ +package runtime + +import ( + "context" + "errors" + "log" + "math/rand/v2" + "time" + + "neo-code/internal/provider" +) + +// handleRunError 负责把运行错误映射为最终 runtime 事件。 +func (s *Service) handleRunError(ctx context.Context, runID string, sessionID string, err error) error { + if errors.Is(err, context.Canceled) { + s.emit(ctx, EventRunCanceled, runID, sessionID, nil) + return context.Canceled + } + + var providerErr *provider.ProviderError + if errors.As(err, &providerErr) { + log.Printf("runtime: provider error (status=%d, code=%s, retryable=%v): %s", + providerErr.StatusCode, providerErr.Code, providerErr.Retryable, providerErr.Message) + } + + s.emit(ctx, EventError, runID, sessionID, err.Error()) + return err +} + +// isRetryableProviderError 判断 provider 错误是否允许 runtime 级重试。 +func isRetryableProviderError(err error) bool { + var providerErr *provider.ProviderError + if !errors.As(err, &providerErr) { + return false + } + return providerErr.Retryable +} + +// providerRetryBackoff 计算 runtime 级 provider 重试等待时间。 +func providerRetryBackoff(attempt int) time.Duration { + wait := providerRetryBaseWait << (attempt - 1) + jitter := float64(wait) * (0.5 + rand.Float64()) + wait = time.Duration(jitter) + if wait > providerRetryMaxWait { + wait = providerRetryMaxWait + } + return wait +} diff --git a/internal/runtime/permission.go b/internal/runtime/permission.go index 4af2ae3c..2cc231c3 100644 --- a/internal/runtime/permission.go +++ b/internal/runtime/permission.go @@ -5,29 +5,42 @@ import ( "errors" "fmt" "strings" - "sync" "time" providertypes "neo-code/internal/provider/types" + approvalflow "neo-code/internal/runtime/approval" "neo-code/internal/security" "neo-code/internal/tools" ) -// PermissionResolutionInput 描述一次来自界面的权限审批决定。 +// PermissionResolutionInput 描述一次来自界面的审批决定。 type PermissionResolutionInput struct { RequestID string Decision PermissionResolutionDecision } -// PermissionResolutionDecision 表示用户在权限提示中的最终选择。 -type PermissionResolutionDecision string +// PermissionResolutionDecision 表示用户在审批弹层中的最终选择。 +type PermissionResolutionDecision = approvalflow.Decision const ( - PermissionResolutionAllowOnce PermissionResolutionDecision = "allow_once" - PermissionResolutionAllowSession PermissionResolutionDecision = "allow_session" - PermissionResolutionReject PermissionResolutionDecision = "reject" + permissionDecisionAllow = "allow" + permissionDecisionDeny = "deny" + + permissionResolvedApproved = "approved" + permissionResolvedRejected = "rejected" + permissionResolvedDenied = "denied" + + permissionReasonApprovedByUser = "permission approved by user" + permissionReasonRejectedByUser = "permission rejected by user" + permissionRejectedErrorMessage = "tools: permission rejected by user" + permissionRejectedContentFmt = "tool error\ntool: %s\nreason: %s" + + permissionToolCategoryFilesystemRead = "filesystem_read" + permissionToolCategoryFilesystemWrite = "filesystem_write" + permissionToolCategoryMCP = "mcp" ) +// permissionExecutionInput 汇总一次工具执行与审批协作所需的上下文。 type permissionExecutionInput struct { RunID string SessionID string @@ -36,63 +49,25 @@ type permissionExecutionInput struct { ToolTimeout time.Duration } -type pendingPermissionRequest struct { - RequestID string - RunID string - SessionID string - Call providertypes.ToolCall - Action security.Action - ResultCh chan PermissionResolutionDecision - Submitted bool -} - -var runtimePendingPermissions = struct { - mu sync.Mutex - nextID uint64 - byRun map[*Service]*pendingPermissionRequest -}{ - byRun: make(map[*Service]*pendingPermissionRequest), -} - -// ResolvePermission 接收 UI 的审批决定，并唤醒 runtime 中等待的工具调用。 +// ResolvePermission 接收 UI 的审批结果，并提交给等待中的工具调用。 func (s *Service) ResolvePermission(ctx context.Context, input PermissionResolutionInput) error { requestID := strings.TrimSpace(input.RequestID) if requestID == "" { return errors.New("runtime: permission request id is empty") } - decision := normalizePermissionResolutionDecision(input.Decision) - if decision == "" { - return fmt.Errorf("runtime: unsupported permission decision %q", input.Decision) - } if err := ctx.Err(); err != nil { return err } - runtimePendingPermissions.mu.Lock() - pending := runtimePendingPermissions.byRun[s] - if pending == nil || pending.RequestID != requestID { - runtimePendingPermissions.mu.Unlock() - return fmt.Errorf("runtime: permission request %q not found", requestID) - } - // Submitted 标记用于避免重复提交同一个 request 导致阻塞。 - if pending.Submitted { - runtimePendingPermissions.mu.Unlock() - return nil - } - pending.Submitted = true - resultCh := pending.ResultCh - runtimePendingPermissions.mu.Unlock() - - // 非阻塞提交，避免 UI 重复触发导致写满 channel 后长时间卡住。 - select { - case resultCh <- decision: - return nil + switch input.Decision { + case approvalflow.DecisionAllowOnce, approvalflow.DecisionAllowSession, approvalflow.DecisionReject: + return s.approvalBroker.Resolve(requestID, input.Decision) default: - return nil + return fmt.Errorf("runtime: unsupported permission decision %q", input.Decision) } } -// executeToolCallWithPermission 执行工具调用并处理 ask 决策的显式审批闭环。 +// executeToolCallWithPermission 执行工具调用，并在 ask/deny 路径上统一发出权限事件。 func (s *Service) executeToolCallWithPermission(ctx context.Context, input permissionExecutionInput) (tools.ToolResult, error) { callInput := tools.ToolCallInput{ ID: input.Call.ID, @@ -120,61 +95,62 @@ func (s *Service) executeToolCallWithPermission(ctx context.Context, input permi return result, execErr } if !strings.EqualFold(permissionErr.Decision(), string(security.DecisionAsk)) { + s.emitPermissionResolved( + ctx, + input, + permissionErr, + "", + permissionErr.Decision(), + permissionResolutionStatus(permissionErr.Decision()), + permissionErr.RememberScope(), + ) return result, execErr } - decision, scope, requestID, err := s.awaitPermissionDecision(ctx, input, permissionErr) + decision, requestID, err := s.awaitPermissionDecision(ctx, input, permissionErr) if err != nil { return result, err } - if decision == PermissionResolutionReject { + scope, err := rememberScopeFromDecision(decision) + if err != nil { + return result, err + } + if decision == approvalflow.DecisionReject { if scope == tools.SessionPermissionScopeReject { if rememberErr := s.toolManager.RememberSessionDecision(input.SessionID, permissionErr.Action(), scope); rememberErr != nil { return tools.ToolResult{}, rememberErr } } - s.emit(ctx, EventPermissionResolved, input.RunID, input.SessionID, PermissionResolvedPayload{ - RequestID: requestID, - ToolCallID: input.Call.ID, - ToolName: input.Call.Name, - ToolCategory: permissionToolCategory(permissionErr.Action()), - ActionType: string(permissionErr.Action().Type), - Operation: permissionErr.Action().Payload.Operation, - TargetType: string(permissionErr.Action().Payload.TargetType), - Target: permissionErr.Action().Payload.Target, - Decision: "deny", - Reason: "permission rejected by user", - RuleID: permissionErr.RuleID(), - RememberScope: string(scope), - ResolvedAs: "rejected", - }) + s.emitPermissionResolved( + ctx, + input, + permissionErr, + requestID, + permissionDecisionDeny, + permissionResolvedRejected, + string(scope), + ) return tools.ToolResult{ ToolCallID: input.Call.ID, Name: input.Call.Name, - Content: "tool error\ntool: " + input.Call.Name + "\nreason: permission rejected by user", + Content: fmt.Sprintf(permissionRejectedContentFmt, input.Call.Name, permissionReasonRejectedByUser), IsError: true, - }, errors.New("tools: permission rejected by user") + }, errors.New(permissionRejectedErrorMessage) } if rememberErr := s.toolManager.RememberSessionDecision(input.SessionID, permissionErr.Action(), scope); rememberErr != nil { return tools.ToolResult{}, rememberErr } - s.emit(ctx, EventPermissionResolved, input.RunID, input.SessionID, PermissionResolvedPayload{ - RequestID: requestID, - ToolCallID: input.Call.ID, - ToolName: input.Call.Name, - ToolCategory: permissionToolCategory(permissionErr.Action()), - ActionType: string(permissionErr.Action().Type), - Operation: permissionErr.Action().Payload.Operation, - TargetType: string(permissionErr.Action().Payload.TargetType), - Target: permissionErr.Action().Payload.Target, - Decision: "allow", - Reason: "permission approved by user", - RuleID: permissionErr.RuleID(), - RememberScope: string(scope), - ResolvedAs: "approved", - }) + s.emitPermissionResolved( + ctx, + input, + permissionErr, + requestID, + permissionDecisionAllow, + permissionResolvedApproved, + string(scope), + ) retryCtx, retryCancel := context.WithTimeout(ctx, input.ToolTimeout) retryResult, retryErr := s.toolManager.Execute(retryCtx, callInput) @@ -182,17 +158,20 @@ func (s *Service) executeToolCallWithPermission(ctx context.Context, input permi return retryResult, retryErr } -// awaitPermissionDecision 发送 permission_request 事件，并等待 UI 回传审批结果。 +// awaitPermissionDecision 发出 permission_request 事件，并等待外部审批结果。 func (s *Service) awaitPermissionDecision( ctx context.Context, input permissionExecutionInput, permissionErr *tools.PermissionDecisionError, -) (PermissionResolutionDecision, tools.SessionPermissionScope, string, error) { - request := registerPendingPermission(s, input, permissionErr.Action()) - defer clearPendingPermission(s, request.RequestID) +) (approvalflow.Decision, string, error) { + requestID, resultCh, err := s.approvalBroker.Open() + if err != nil { + return "", "", err + } + defer s.approvalBroker.Close(requestID) s.emit(ctx, EventPermissionRequest, input.RunID, input.SessionID, PermissionRequestPayload{ - RequestID: request.RequestID, + RequestID: requestID, ToolCallID: input.Call.ID, ToolName: input.Call.Name, ToolCategory: permissionToolCategory(permissionErr.Action()), @@ -207,91 +186,92 @@ func (s *Service) awaitPermissionDecision( select { case <-ctx.Done(): - return "", "", request.RequestID, ctx.Err() - case decision := <-request.ResultCh: - scope, err := rememberScopeFromDecision(decision) - if err != nil { - return "", "", request.RequestID, err - } - return decision, scope, request.RequestID, nil + return "", requestID, ctx.Err() + case decision := <-resultCh: + return decision, requestID, nil } } -// registerPendingPermission 注册待审批请求并返回 request id。 -func registerPendingPermission(s *Service, input permissionExecutionInput, action security.Action) pendingPermissionRequest { - runtimePendingPermissions.mu.Lock() - defer runtimePendingPermissions.mu.Unlock() - - runtimePendingPermissions.nextID++ - request := pendingPermissionRequest{ - RequestID: fmt.Sprintf("perm-%d", runtimePendingPermissions.nextID), - RunID: input.RunID, - SessionID: input.SessionID, - Call: input.Call, - Action: action, - ResultCh: make(chan PermissionResolutionDecision, 1), +// emitPermissionResolved 将权限处理结果统一映射为 runtime 事件。 +func (s *Service) emitPermissionResolved( + ctx context.Context, + input permissionExecutionInput, + permissionErr *tools.PermissionDecisionError, + requestID string, + decision string, + resolvedAs string, + rememberScope string, +) { + reason := strings.TrimSpace(permissionErr.Reason()) + if strings.TrimSpace(requestID) != "" { + switch resolvedAs { + case permissionResolvedApproved: + reason = permissionReasonApprovedByUser + case permissionResolvedRejected: + reason = permissionReasonRejectedByUser + } } - runtimePendingPermissions.byRun[s] = &request - return request -} - -// clearPendingPermission 清理待审批请求，避免跨请求误用。 -func clearPendingPermission(s *Service, requestID string) { - runtimePendingPermissions.mu.Lock() - defer runtimePendingPermissions.mu.Unlock() - - current := runtimePendingPermissions.byRun[s] - if current != nil && current.RequestID == requestID { - delete(runtimePendingPermissions.byRun, s) + if reason == "" { + reason = strings.TrimSpace(permissionErr.Error()) } -} -// normalizePermissionResolutionDecision 将审批决定归一化为受支持枚举值。 -func normalizePermissionResolutionDecision(decision PermissionResolutionDecision) PermissionResolutionDecision { - switch strings.ToLower(strings.TrimSpace(string(decision))) { - case "allow_once", "once", "y", "yes": - return PermissionResolutionAllowOnce - case "allow_session", "always", "always_session", "a": - return PermissionResolutionAllowSession - case "reject", "deny", "n", "no", "r": - return PermissionResolutionReject - default: - return "" - } + s.emit(ctx, EventPermissionResolved, input.RunID, input.SessionID, PermissionResolvedPayload{ + RequestID: requestID, + ToolCallID: input.Call.ID, + ToolName: input.Call.Name, + ToolCategory: permissionToolCategory(permissionErr.Action()), + ActionType: string(permissionErr.Action().Type), + Operation: permissionErr.Action().Payload.Operation, + TargetType: string(permissionErr.Action().Payload.TargetType), + Target: permissionErr.Action().Payload.Target, + Decision: decision, + Reason: reason, + RuleID: permissionErr.RuleID(), + RememberScope: rememberScope, + ResolvedAs: resolvedAs, + }) } -// rememberScopeFromDecision 将审批决定映射到工具层权限记忆范围。 +// rememberScopeFromDecision 将审批结果映射为工具层的会话记忆范围。 func rememberScopeFromDecision(decision PermissionResolutionDecision) (tools.SessionPermissionScope, error) { switch decision { - case PermissionResolutionAllowOnce: + case approvalflow.DecisionAllowOnce: return tools.SessionPermissionScopeOnce, nil - case PermissionResolutionAllowSession: + case approvalflow.DecisionAllowSession: return tools.SessionPermissionScopeAlways, nil - case PermissionResolutionReject: + case approvalflow.DecisionReject: return tools.SessionPermissionScopeReject, nil default: return "", fmt.Errorf("runtime: unsupported permission decision %q", decision) } } -// permissionToolCategory 将 action 归一化为工具类别标签，供审批展示和记忆使用。 +// permissionResolutionStatus 负责将工具层决策映射为 resolved_as 字段。 +func permissionResolutionStatus(decision string) string { + if strings.EqualFold(strings.TrimSpace(decision), string(security.DecisionAsk)) { + return permissionResolvedRejected + } + return permissionResolvedDenied +} + +// permissionToolCategory 将安全动作收敛为用于审批展示的工具分类标签。 func permissionToolCategory(action security.Action) string { resource := strings.ToLower(strings.TrimSpace(action.Payload.Resource)) switch action.Type { case security.ActionTypeRead: if strings.HasPrefix(resource, "filesystem_") { - return "filesystem_read" + return permissionToolCategoryFilesystemRead } case security.ActionTypeWrite: if strings.HasPrefix(resource, "filesystem_") { - return "filesystem_write" + return permissionToolCategoryFilesystemWrite } case security.ActionTypeMCP: target := strings.ToLower(strings.TrimSpace(action.Payload.Target)) if serverIdentity := security.CanonicalMCPServerIdentity(target); serverIdentity != "" { return serverIdentity } - return "mcp" + return permissionToolCategoryMCP } if resource != "" { return resource diff --git a/internal/runtime/permission_test.go b/internal/runtime/permission_test.go index 6e00dabe..9895f4da 100644 --- a/internal/runtime/permission_test.go +++ b/internal/runtime/permission_test.go @@ -8,6 +8,7 @@ import ( "time" providertypes "neo-code/internal/provider/types" + approvalflow "neo-code/internal/runtime/approval" "neo-code/internal/security" "neo-code/internal/tools" ) @@ -34,7 +35,7 @@ func TestResolvePermissionValidation(t *testing.T) { } if err := service.ResolvePermission(context.Background(), PermissionResolutionInput{ RequestID: "perm-not-found", - Decision: PermissionResolutionAllowOnce, + Decision: approvalflow.DecisionAllowOnce, }); err == nil { t.Fatalf("expected request not found error") } @@ -51,36 +52,23 @@ func TestResolvePermissionSuccess(t *testing.T) { nil, ) - request := registerPendingPermission(service, permissionExecutionInput{ - RunID: "run-permission", - SessionID: "session-permission", - Call: providertypes.ToolCall{ - ID: "call-1", - Name: "webfetch", - }, - }, security.Action{ - Type: security.ActionTypeRead, - Payload: security.ActionPayload{ - ToolName: "webfetch", - Resource: "webfetch", - Operation: "fetch", - TargetType: security.TargetTypeURL, - Target: "https://example.com", - }, - }) - defer clearPendingPermission(service, request.RequestID) + requestID, resultCh, err := service.approvalBroker.Open() + if err != nil { + t.Fatalf("open approval request: %v", err) + } + defer service.approvalBroker.Close(requestID) errCh := make(chan error, 1) go func() { errCh <- service.ResolvePermission(context.Background(), PermissionResolutionInput{ - RequestID: request.RequestID, - Decision: PermissionResolutionAllowSession, + RequestID: requestID, + Decision: approvalflow.DecisionAllowSession, }) }() select { - case resolved := <-request.ResultCh: - if resolved != PermissionResolutionAllowSession { + case resolved := <-resultCh: + if resolved != approvalflow.DecisionAllowSession { t.Fatalf("expected allow session decision, got %q", resolved) } case <-time.After(2 * time.Second): @@ -103,25 +91,15 @@ func TestResolvePermissionDuplicateSubmissionIsNonBlocking(t *testing.T) { nil, ) - request := registerPendingPermission(service, permissionExecutionInput{ - RunID: "run-permission-dup", - SessionID: "session-permission-dup", - Call: providertypes.ToolCall{ - ID: "call-dup", - Name: "webfetch", - }, - }, security.Action{ - Type: security.ActionTypeRead, - Payload: security.ActionPayload{ - ToolName: "webfetch", - Resource: "webfetch", - }, - }) - defer clearPendingPermission(service, request.RequestID) + requestID, _, err := service.approvalBroker.Open() + if err != nil { + t.Fatalf("open approval request: %v", err) + } + defer service.approvalBroker.Close(requestID) if err := service.ResolvePermission(context.Background(), PermissionResolutionInput{ - RequestID: request.RequestID, - Decision: PermissionResolutionAllowOnce, + RequestID: requestID, + Decision: approvalflow.DecisionAllowOnce, }); err != nil { t.Fatalf("first ResolvePermission() error = %v", err) } @@ -129,8 +107,8 @@ func TestResolvePermissionDuplicateSubmissionIsNonBlocking(t *testing.T) { secondDone := make(chan error, 1) go func() { secondDone <- service.ResolvePermission(context.Background(), PermissionResolutionInput{ - RequestID: request.RequestID, - Decision: PermissionResolutionAllowSession, + RequestID: requestID, + Decision: approvalflow.DecisionAllowSession, }) }() @@ -215,7 +193,7 @@ waitRequest: if err := service.ResolvePermission(context.Background(), PermissionResolutionInput{ RequestID: requestPayload.RequestID, - Decision: PermissionResolutionReject, + Decision: approvalflow.DecisionReject, }); err != nil { t.Fatalf("ResolvePermission() error = %v", err) } @@ -251,26 +229,13 @@ waitRequest: func TestPermissionHelpers(t *testing.T) { t.Parallel() - if got := normalizePermissionResolutionDecision(PermissionResolutionDecision("Y")); got != PermissionResolutionAllowOnce { - t.Fatalf("expected Y => allow_once, got %q", got) - } - if got := normalizePermissionResolutionDecision(PermissionResolutionDecision("a")); got != PermissionResolutionAllowSession { - t.Fatalf("expected a => allow_session, got %q", got) - } - if got := normalizePermissionResolutionDecision(PermissionResolutionDecision("n")); got != PermissionResolutionReject { - t.Fatalf("expected n => reject, got %q", got) - } - if got := normalizePermissionResolutionDecision(PermissionResolutionDecision("???")); got != "" { - t.Fatalf("expected unknown => empty, got %q", got) - } - - if scope, err := rememberScopeFromDecision(PermissionResolutionAllowOnce); err != nil || scope != tools.SessionPermissionScopeOnce { + if scope, err := rememberScopeFromDecision(approvalflow.DecisionAllowOnce); err != nil || scope != tools.SessionPermissionScopeOnce { t.Fatalf("expected once scope, got %q / %v", scope, err) } - if scope, err := rememberScopeFromDecision(PermissionResolutionAllowSession); err != nil || scope != tools.SessionPermissionScopeAlways { + if scope, err := rememberScopeFromDecision(approvalflow.DecisionAllowSession); err != nil || scope != tools.SessionPermissionScopeAlways { t.Fatalf("expected always scope, got %q / %v", scope, err) } - if scope, err := rememberScopeFromDecision(PermissionResolutionReject); err != nil || scope != tools.SessionPermissionScopeReject { + if scope, err := rememberScopeFromDecision(approvalflow.DecisionReject); err != nil || scope != tools.SessionPermissionScopeReject { t.Fatalf("expected reject scope, got %q / %v", scope, err) } if _, err := rememberScopeFromDecision(PermissionResolutionDecision("invalid")); err == nil { @@ -284,8 +249,8 @@ func TestPermissionHelpers(t *testing.T) { Resource: "filesystem_grep", }, }) - if category != "filesystem_read" { - t.Fatalf("expected filesystem_read category, got %q", category) + if category != permissionToolCategoryFilesystemRead { + t.Fatalf("expected %s category, got %q", permissionToolCategoryFilesystemRead, category) } category = permissionToolCategory(security.Action{ @@ -315,8 +280,8 @@ func TestPermissionHelpers(t *testing.T) { Target: "mcp", }, }) - if category != "mcp" { - t.Fatalf("expected mcp fallback category, got %q", category) + if category != permissionToolCategoryMCP { + t.Fatalf("expected %s fallback category, got %q", permissionToolCategoryMCP, category) } } @@ -330,28 +295,18 @@ func TestResolvePermissionCanceledContext(t *testing.T) { &scriptedProviderFactory{provider: &scriptedProvider{}}, nil, ) - request := registerPendingPermission(service, permissionExecutionInput{ - RunID: "run-canceled", - SessionID: "session-canceled", - Call: providertypes.ToolCall{ - ID: "call-canceled", - Name: "webfetch", - }, - }, security.Action{ - Type: security.ActionTypeRead, - Payload: security.ActionPayload{ - ToolName: "webfetch", - Resource: "webfetch", - }, - }) - defer clearPendingPermission(service, request.RequestID) - request.ResultCh <- PermissionResolutionAllowOnce + requestID, resultCh, err := service.approvalBroker.Open() + if err != nil { + t.Fatalf("open approval request: %v", err) + } + defer service.approvalBroker.Close(requestID) + resultCh <- approvalflow.DecisionAllowOnce ctx, cancel := context.WithCancel(context.Background()) cancel() if err := service.ResolvePermission(ctx, PermissionResolutionInput{ - RequestID: request.RequestID, - Decision: PermissionResolutionAllowOnce, + RequestID: requestID, + Decision: approvalflow.DecisionAllowOnce, }); !errors.Is(err, context.Canceled) { t.Fatalf("expected context canceled, got %v", err) } diff --git a/internal/runtime/provider_stream.go b/internal/runtime/provider_stream.go new file mode 100644 index 00000000..97dd5992 --- /dev/null +++ b/internal/runtime/provider_stream.go @@ -0,0 +1,89 @@ +package runtime + +import ( + "context" + "fmt" + + "neo-code/internal/provider" + "neo-code/internal/provider/streaming" + providertypes "neo-code/internal/provider/types" +) + +// streamGenerateResult 统一承载一次流式生成的消息、用量与消费错误。 +type streamGenerateResult struct { + message providertypes.Message + inputTokens int + outputTokens int + err error +} + +// generateStreamingMessage 负责执行一次基于流式事件的生成调用，并收敛最终 assistant 消息与 usage。 +func generateStreamingMessage( + ctx context.Context, + modelProvider provider.Provider, + req providertypes.GenerateRequest, + hooks streaming.Hooks, +) streamGenerateResult { + acc := streaming.NewAccumulator() + streamEvents := make(chan providertypes.StreamEvent, 32) + streamDone := make(chan streamGenerateResult, 1) + + go func() { + outcome := streamGenerateResult{} + defer func() { + streamDone <- outcome + }() + + for { + select { + case event, ok := <-streamEvents: + if !ok { + return + } + if err := streaming.HandleEvent(event, acc, streaming.Hooks{ + OnTextDelta: hooks.OnTextDelta, + OnToolCallStart: hooks.OnToolCallStart, + OnMessageDone: func(payload providertypes.MessageDonePayload) { + if payload.Usage != nil { + outcome.inputTokens = payload.Usage.InputTokens + outcome.outputTokens = payload.Usage.OutputTokens + } + if hooks.OnMessageDone != nil { + hooks.OnMessageDone(payload) + } + }, + }); err != nil && outcome.err == nil { + outcome.err = err + } + case <-ctx.Done(): + return + } + } + }() + + generateErr := modelProvider.Generate(ctx, req, streamEvents) + close(streamEvents) + outcome := <-streamDone + if outcome.err != nil { + if generateErr != nil { + outcome.err = fmt.Errorf("runtime: provider stream handling failed after provider error: %v: %w", generateErr, outcome.err) + } + return outcome + } + if generateErr != nil { + outcome.err = generateErr + return outcome + } + if !acc.MessageDone() { + outcome.err = fmt.Errorf("%w: provider stream ended without message_done event", provider.ErrStreamInterrupted) + return outcome + } + + message, err := acc.BuildMessage() + if err != nil { + outcome.err = err + return outcome + } + outcome.message = message + return outcome +} diff --git a/internal/runtime/run.go b/internal/runtime/run.go new file mode 100644 index 00000000..0821fb9c --- /dev/null +++ b/internal/runtime/run.go @@ -0,0 +1,272 @@ +package runtime + +import ( + "context" + "errors" + "fmt" + "strings" + "time" + + "neo-code/internal/config" + agentcontext "neo-code/internal/context" + contextcompact "neo-code/internal/context/compact" + "neo-code/internal/provider" + "neo-code/internal/provider/streaming" + providertypes "neo-code/internal/provider/types" + agentsession "neo-code/internal/session" + "neo-code/internal/tools" +) + +// Run 执行一次完整的 ReAct 闭环：保存用户输入、驱动模型、执行工具并发出事件。 +// 会话创建/加载在锁外完成，仅会话级操作（ReAct 循环）受会话锁保护， +// 不同会话可并行执行。 +func (s *Service) Run(ctx context.Context, input UserInput) error { + if strings.TrimSpace(input.Content) == "" { + return errors.New("runtime: input content is empty") + } + + runCtx, cancel := context.WithCancel(ctx) + runToken := s.startRun(cancel) + defer func() { + cancel() + s.finishRun(runToken) + }() + ctx = runCtx + + initialCfg := s.configManager.Get() + session, err := s.loadOrCreateSession(ctx, input.SessionID, input.Content, initialCfg.Workdir, input.Workdir) + if err != nil { + return s.handleRunError(ctx, input.RunID, input.SessionID, err) + } + + sessionMu := s.acquireSessionLock(session.ID) + sessionMu.Lock() + defer sessionMu.Unlock() + + state := newRunState(input.RunID, session) + if err := s.appendUserMessageAndSave(ctx, &state, input.Content); err != nil { + return s.handleRunError(ctx, state.runID, state.session.ID, err) + } + + for turn := 0; ; turn++ { + maxLoops := resolveMaxLoops(s.configManager.Get()) + if turn >= maxLoops { + err := errors.New("runtime: max loop reached") + s.emit(ctx, EventError, state.runID, state.session.ID, err.Error()) + return err + } + + for { + if err := ctx.Err(); err != nil { + return s.handleRunError(ctx, state.runID, state.session.ID, err) + } + + snapshot, rebuilt, err := s.prepareTurnSnapshot(ctx, &state) + if err != nil { + return s.handleRunError(ctx, state.runID, state.session.ID, err) + } + if rebuilt { + continue + } + + turnResult, err := s.callProviderWithRetry(ctx, &state, snapshot) + if err != nil { + if provider.IsContextTooLong(err) && !state.reactiveCompactUsed { + state.reactiveCompactUsed = true + _, _ = s.applyCompactForState(ctx, &state, snapshot.config, contextcompact.ModeReactive, compactErrorBestEffort) + continue + } + return s.handleRunError(ctx, state.runID, state.session.ID, err) + } + + if strings.TrimSpace(turnResult.assistant.Role) == "" { + turnResult.assistant.Role = providertypes.RoleAssistant + } + state.recordUsage(turnResult.inputTokens, turnResult.outputTokens) + if err := s.appendAssistantMessageAndSave(ctx, &state, snapshot, turnResult.assistant); err != nil { + return s.handleRunError(ctx, state.runID, state.session.ID, err) + } + s.emitTokenUsage(ctx, &state, turnResult) + + if len(turnResult.assistant.ToolCalls) == 0 { + s.emit(ctx, EventAgentDone, state.runID, state.session.ID, turnResult.assistant) + return nil + } + if err := s.executeAssistantToolCalls(ctx, &state, snapshot, turnResult.assistant); err != nil { + return s.handleRunError(ctx, state.runID, state.session.ID, err) + } + break + } + } +} + +// prepareTurnSnapshot 基于当前会话状态冻结一轮推理所需的请求快照。 +func (s *Service) prepareTurnSnapshot(ctx context.Context, state *runState) (turnSnapshot, bool, error) { + cfg := s.configManager.Get() + activeWorkdir := agentsession.EffectiveWorkdir(state.session.Workdir, cfg.Workdir) + + builtContext, err := s.contextBuilder.Build(ctx, agentcontext.BuildInput{ + Messages: state.session.Messages, + Metadata: agentcontext.Metadata{ + Workdir: activeWorkdir, + Shell: cfg.Shell, + Provider: cfg.SelectedProvider, + Model: cfg.CurrentModel, + SessionInputTokens: state.tokenInputTotal, + SessionOutputTokens: state.tokenOutputTotal, + }, + Compact: agentcontext.CompactOptions{ + DisableMicroCompact: cfg.Context.Compact.MicroCompactDisabled, + AutoCompactThreshold: autoCompactThreshold(cfg), + }, + }) + if err != nil { + return turnSnapshot{}, false, err + } + + if builtContext.AutoCompactSuggested && !state.compactApplied { + applied, err := s.applyCompactForState(ctx, state, cfg, contextcompact.ModeAuto, compactErrorBestEffort) + if err != nil { + return turnSnapshot{}, false, err + } + if applied { + return turnSnapshot{}, true, nil + } + } + + toolSpecs, err := s.toolManager.ListAvailableSpecs(ctx, tools.SpecListInput{ + SessionID: state.session.ID, + }) + if err != nil { + return turnSnapshot{}, false, err + } + + resolvedProvider, err := config.ResolveSelectedProvider(cfg) + if err != nil { + return turnSnapshot{}, false, err + } + + model := strings.TrimSpace(cfg.CurrentModel) + return turnSnapshot{ + config: cfg, + providerConfig: resolvedProvider.ToRuntimeConfig(), + model: model, + workdir: activeWorkdir, + toolTimeout: time.Duration(cfg.ToolTimeoutSec) * time.Second, + request: providertypes.GenerateRequest{ + Model: model, + SystemPrompt: builtContext.SystemPrompt, + Messages: builtContext.Messages, + Tools: toolSpecs, + }, + }, false, nil +} + +// callProviderWithRetry 使用冻结后的 turnSnapshot 执行 provider 调用与必要重试。 +func (s *Service) callProviderWithRetry( + ctx context.Context, + state *runState, + snapshot turnSnapshot, +) (providerTurnResult, error) { + var lastErr error + + for retryAttempt := 0; retryAttempt <= defaultProviderRetryMax; retryAttempt++ { + if retryAttempt > 0 { + wait := providerRetryBackoff(retryAttempt) + s.emit(ctx, EventProviderRetry, state.runID, state.session.ID, + fmt.Sprintf("retrying provider call (attempt %d/%d, wait=%.1fs)...", + retryAttempt, defaultProviderRetryMax, wait.Seconds())) + + select { + case <-ctx.Done(): + return providerTurnResult{}, ctx.Err() + case <-time.After(wait): + } + } + + modelProvider, err := s.providerFactory.Build(ctx, snapshot.providerConfig) + if err != nil { + return providerTurnResult{}, err + } + + streamOutcome := generateStreamingMessage(ctx, modelProvider, snapshot.request, streaming.Hooks{ + OnTextDelta: func(text string) { + s.emit(ctx, EventAgentChunk, state.runID, state.session.ID, text) + }, + OnToolCallStart: func(payload providertypes.ToolCallStartPayload) { + s.emit(ctx, EventToolCallThinking, state.runID, state.session.ID, payload.Name) + }, + }) + if streamOutcome.err != nil { + lastErr = streamOutcome.err + if !isRetryableProviderError(lastErr) { + return providerTurnResult{}, lastErr + } + if ctx.Err() != nil { + return providerTurnResult{}, ctx.Err() + } + continue + } + + return providerTurnResult{ + assistant: streamOutcome.message, + inputTokens: streamOutcome.inputTokens, + outputTokens: streamOutcome.outputTokens, + }, nil + } + + if lastErr == nil { + lastErr = errors.New("max retries exceeded") + } + return providerTurnResult{}, fmt.Errorf("runtime: max retries exhausted, last error: %w", lastErr) +} + +// emitTokenUsage 在单轮 provider 调用成功后发出 token_usage 事件。 +func (s *Service) emitTokenUsage(ctx context.Context, state *runState, result providerTurnResult) { + if result.inputTokens == 0 && result.outputTokens == 0 { + return + } + s.emit(ctx, EventTokenUsage, state.runID, state.session.ID, TokenUsagePayload{ + InputTokens: result.inputTokens, + OutputTokens: result.outputTokens, + SessionInputTokens: state.tokenInputTotal, + SessionOutputTokens: state.tokenOutputTotal, + }) +} + +// applyCompactForState 在运行中执行 compact，并把结果同步回 runState。 +func (s *Service) applyCompactForState( + ctx context.Context, + state *runState, + cfg config.Config, + mode contextcompact.Mode, + errorPolicy compactErrorPolicy, +) (bool, error) { + session, result, err := s.runCompactForSession(ctx, state.runID, state.session, cfg, mode, errorPolicy) + if err != nil { + return false, err + } + state.session = session + if result.Applied { + state.resetTokenTotals() + state.compactApplied = true + return true, nil + } + return false, nil +} + +// resolveMaxLoops 收敛运行时最大推理轮数的默认值逻辑。 +func resolveMaxLoops(cfg config.Config) int { + if cfg.MaxLoops <= 0 { + return defaultMaxLoops + } + return cfg.MaxLoops +} + +// autoCompactThreshold 返回当前配置下的自动 compact 触发阈值。 +func autoCompactThreshold(cfg config.Config) int { + if cfg.Context.AutoCompact.Enabled && cfg.Context.AutoCompact.InputTokenThreshold > 0 { + return cfg.Context.AutoCompact.InputTokenThreshold + } + return 0 +} diff --git a/internal/runtime/runtime.go b/internal/runtime/runtime.go index 60be921b..d5738da0 100644 --- a/internal/runtime/runtime.go +++ b/internal/runtime/runtime.go @@ -2,12 +2,7 @@ package runtime import ( "context" - "errors" - "fmt" - "log" - "math/rand/v2" "path/filepath" - "sort" "strings" "sync" "time" @@ -16,99 +11,19 @@ import ( agentcontext "neo-code/internal/context" contextcompact "neo-code/internal/context/compact" "neo-code/internal/provider" - providertypes "neo-code/internal/provider/types" + "neo-code/internal/runtime/approval" agentsession "neo-code/internal/session" "neo-code/internal/tools" ) const ( - // defaultProviderRetryMax 是 runtime 层对单次 provider.Generate() 的最大重试次数（不含首次调用）。 - // 与 RetryTransport 的 HTTP 层重试互补：Transport 耗尽后 runtime 仍可重试整个 Chat 调用。 defaultProviderRetryMax = 2 - // providerRetryBaseWait 是 runtime 层重试的初始等待时间。 - providerRetryBaseWait = 1 * time.Second - // providerRetryMaxWait 是 runtime 层重试的最大等待时间。 - providerRetryMaxWait = 5 * time.Second + providerRetryBaseWait = 1 * time.Second + providerRetryMaxWait = 5 * time.Second + defaultMaxLoops = 8 ) -// streamAccumulator 在流式事件处理过程中累积本轮对话需要持久化的助手消息状态， -// 包括文本内容和工具调用列表。 -type streamAccumulator struct { - content strings.Builder - toolCalls map[int]*providertypes.ToolCall - messageDone bool -} - -// newStreamAccumulator 创建并初始化一个空的流式事件累积器。 -func newStreamAccumulator() *streamAccumulator { - return &streamAccumulator{ - toolCalls: make(map[int]*providertypes.ToolCall), - } -} - -// accumulateTextDelta 累积文本增量片段。 -func (a *streamAccumulator) accumulateTextDelta(text string) { - a.content.WriteString(text) -} - -// ensureToolCall 返回指定索引的工具调用条目，不存在时会先创建占位对象。 -func (a *streamAccumulator) ensureToolCall(index int) *providertypes.ToolCall { - call, exists := a.toolCalls[index] - if !exists { - call = &providertypes.ToolCall{} - a.toolCalls[index] = call - } - return call -} - -// accumulateToolCallStart 记录新发现的工具调用（首次出现时创建条目）。 -func (a *streamAccumulator) accumulateToolCallStart(index int, id, name string) { - call := a.ensureToolCall(index) - if strings.TrimSpace(id) != "" { - call.ID = id - } - if strings.TrimSpace(name) != "" { - call.Name = name - } -} - -// accumulateToolCallDelta 累积工具调用参数增量。 -func (a *streamAccumulator) accumulateToolCallDelta(index int, id, argumentsDelta string) { - call := a.ensureToolCall(index) - if strings.TrimSpace(id) != "" { - call.ID = id - } - call.Arguments += argumentsDelta -} - -// buildMessage 从累积状态构建最终的 assistant Message 对象，并校验工具调用元数据是否完整。 -func (a *streamAccumulator) buildMessage() (providertypes.Message, error) { - ordered := make([]int, 0, len(a.toolCalls)) - for index := range a.toolCalls { - ordered = append(ordered, index) - } - sort.Ints(ordered) - - message := providertypes.Message{ - Role: providertypes.RoleAssistant, - Content: a.content.String(), - } - for _, index := range ordered { - call := a.toolCalls[index] - if call == nil { - continue - } - if strings.TrimSpace(call.ID) == "" { - return providertypes.Message{}, fmt.Errorf("runtime: provider emitted tool call %d without id", index) - } - if strings.TrimSpace(call.Name) == "" { - return providertypes.Message{}, fmt.Errorf("runtime: provider emitted tool call %d without name", index) - } - message.ToolCalls = append(message.ToolCalls, *call) - } - return message, nil -} - +// Runtime 定义 runtime 对外暴露的运行、压缩与审批接口。 type Runtime interface { Run(ctx context.Context, input UserInput) error Compact(ctx context.Context, input CompactInput) (CompactResult, error) @@ -119,6 +34,7 @@ type Runtime interface { LoadSession(ctx context.Context, id string) (agentsession.Session, error) } +// UserInput 描述一次用户输入请求的最小运行参数。 type UserInput struct { SessionID string RunID string @@ -126,27 +42,31 @@ type UserInput struct { Workdir string } +// ProviderFactory 负责基于运行期配置创建 provider 实例。 type ProviderFactory interface { Build(ctx context.Context, cfg provider.RuntimeConfig) (provider.Provider, error) } +// Service 是 runtime 的默认实现，负责组织一次完整的 agent 运行闭环。 type Service struct { - configManager *config.Manager // 配置管理器，提供当前选中的 provider、model、workdir 等配置读取能力。 - sessionStore agentsession.Store // 会话持久化接口，负责保存和加载聊天会话。 - toolManager tools.Manager // 工具管理器，统一工具 schema 暴露与执行入口。 - providerFactory ProviderFactory // Provider 工厂接口，根据配置动态创建具体的 provider 实例。 - contextBuilder agentcontext.Builder // 上下文构建器，负责组装 system prompt 与本轮发给模型的消息上下文。 - compactRunner contextcompact.Runner - events chan RuntimeEvent - operationMu sync.Mutex // 运行级互斥：串行化 Run 与 Compact，避免并发写同一会话。 - runMu sync.Mutex // 仅保护 activeRun* 字段的并发读写。 - activeRunToken uint64 // 当前活跃运行的令牌标识，用于标记正在执行的 Run 实例。 - nextRunToken uint64 // 下一个运行令牌的递增计数器，用于区分不同 Run 的生命周期。 - activeRunCancel context.CancelFunc // 当前活跃 Run 的取消函数。 - sessionInputTokens int // 当前会话累计输入 token。 - sessionOutputTokens int // 当前会话累计输出 token。 -} - + configManager *config.Manager + sessionStore agentsession.Store + toolManager tools.Manager + providerFactory ProviderFactory + contextBuilder agentcontext.Builder + compactRunner contextcompact.Runner + approvalBroker *approval.Broker + + events chan RuntimeEvent + sessionMu sync.Mutex + sessionLocks map[string]*sync.Mutex + runMu sync.Mutex + activeRunToken uint64 + nextRunToken uint64 + activeRunCancel context.CancelFunc +} + +// NewWithFactory 使用注入依赖构建默认 runtime Service。 func NewWithFactory( configManager *config.Manager, toolManager tools.Manager, @@ -170,229 +90,13 @@ func NewWithFactory( toolManager: toolManager, providerFactory: providerFactory, contextBuilder: contextBuilder, + approvalBroker: approval.NewBroker(), events: make(chan RuntimeEvent, 128), + sessionLocks: make(map[string]*sync.Mutex), } } -func (s *Service) Run(ctx context.Context, input UserInput) error { - s.operationMu.Lock() - defer s.operationMu.Unlock() - - runCtx, cancel := context.WithCancel(ctx) - runToken := s.startRun(cancel) - defer func() { - cancel() - s.finishRun(runToken) - }() - ctx = runCtx - - if strings.TrimSpace(input.Content) == "" { - return errors.New("runtime: input content is empty") - } - - initialCfg := s.configManager.Get() - session, err := s.loadOrCreateSession(ctx, input.SessionID, input.Content, initialCfg.Workdir, input.Workdir) - if err != nil { - return s.handleRunError(ctx, input.RunID, input.SessionID, err) - } - s.restoreSessionTokens(session) - - userMessage := providertypes.Message{ - Role: providertypes.RoleUser, - Content: input.Content, - } - session.Messages = append(session.Messages, userMessage) - session.UpdatedAt = time.Now() - if err := s.sessionStore.Save(ctx, &session); err != nil { - return s.handleRunError(ctx, input.RunID, session.ID, err) - } - s.emit(ctx, EventUserMessage, input.RunID, session.ID, userMessage) - - autoCompacted := false - reactiveRetried := false - - for attempt := 0; ; attempt++ { - if err := ctx.Err(); err != nil { - return s.handleRunError(ctx, input.RunID, session.ID, err) - } - - cfg := s.configManager.Get() - activeWorkdir := effectiveSessionWorkdir(session.Workdir, cfg.Workdir) - maxLoops := cfg.MaxLoops - if maxLoops <= 0 { - maxLoops = 8 - } - if attempt >= maxLoops { - err := errors.New("runtime: max loop reached") - s.emit(ctx, EventError, input.RunID, session.ID, err.Error()) - return err - } - - builtContext, err := s.contextBuilder.Build(ctx, agentcontext.BuildInput{ - Messages: session.Messages, - Metadata: agentcontext.Metadata{ - Workdir: activeWorkdir, - Shell: cfg.Shell, - Provider: cfg.SelectedProvider, - Model: cfg.CurrentModel, - SessionInputTokens: s.sessionInputTokens, - SessionOutputTokens: s.sessionOutputTokens, - }, - Compact: agentcontext.CompactOptions{ - DisableMicroCompact: cfg.Context.Compact.MicroCompactDisabled, - AutoCompactThreshold: s.autoCompactThreshold(cfg), - }, - }) - if err != nil { - return s.handleRunError(ctx, input.RunID, session.ID, err) - } - - if builtContext.ShouldAutoCompact && !autoCompacted { - var compactResult contextcompact.Result - session, compactResult, _ = s.runCompactForSession(ctx, input.RunID, session, cfg, contextcompact.ModeAuto, false) - if compactResult.Applied { - autoCompacted = true - s.resetSessionTokenTotals(&session) - // 自动 compact 成功后需要在同一轮重建上下文，避免继续沿用压缩前的请求内容。 - attempt-- - continue - } - } - - toolSpecs, err := s.toolManager.ListAvailableSpecs(ctx, tools.SpecListInput{ - SessionID: session.ID, - }) - if err != nil { - return s.handleRunError(ctx, input.RunID, session.ID, err) - } - - acc, err := s.callProviderWithRetry(ctx, input.RunID, session.ID, providertypes.GenerateRequest{ - Model: cfg.CurrentModel, - SystemPrompt: builtContext.SystemPrompt, - Messages: builtContext.Messages, - Tools: toolSpecs, - }) - if err != nil { - if provider.IsContextTooLong(err) && !reactiveRetried { - reactiveRetried = true - var compactResult contextcompact.Result - session, compactResult, _ = s.runCompactForSession( - ctx, - input.RunID, - session, - cfg, - contextcompact.ModeReactive, - false, - ) - if compactResult.Applied { - s.resetSessionTokenTotals(&session) - autoCompacted = true - } - // Don't count the reactive-compact retry as a new reasoning turn. - attempt-- - continue - } - return s.handleRunError(ctx, input.RunID, session.ID, err) - } - if err := ctx.Err(); err != nil { - return s.handleRunError(ctx, input.RunID, session.ID, err) - } - - metadataChanged := session.Provider != cfg.SelectedProvider || session.Model != cfg.CurrentModel - session.Provider = cfg.SelectedProvider - session.Model = cfg.CurrentModel - session.TokenInputTotal = s.sessionInputTokens - session.TokenOutputTotal = s.sessionOutputTokens - - assistant, err := acc.buildMessage() - if err != nil { - return s.handleRunError(ctx, input.RunID, session.ID, err) - } - if strings.TrimSpace(assistant.Role) == "" { - assistant.Role = providertypes.RoleAssistant - } - - if strings.TrimSpace(assistant.Content) != "" || len(assistant.ToolCalls) > 0 { - session.Messages = append(session.Messages, assistant) - session.UpdatedAt = time.Now() - if err := s.sessionStore.Save(ctx, &session); err != nil { - return s.handleRunError(ctx, input.RunID, session.ID, err) - } - } else if metadataChanged { - session.UpdatedAt = time.Now() - if err := s.sessionStore.Save(ctx, &session); err != nil { - return s.handleRunError(ctx, input.RunID, session.ID, err) - } - } - - if err := ctx.Err(); err != nil { - return s.handleRunError(ctx, input.RunID, session.ID, err) - } - if len(assistant.ToolCalls) == 0 { - s.emit(ctx, EventAgentDone, input.RunID, session.ID, assistant) - return nil - } - - for _, call := range assistant.ToolCalls { - if err := ctx.Err(); err != nil { - return s.handleRunError(ctx, input.RunID, session.ID, err) - } - s.emit(ctx, EventToolStart, input.RunID, session.ID, call) - - result, execErr := s.executeToolCallWithPermission(ctx, permissionExecutionInput{ - RunID: input.RunID, - SessionID: session.ID, - Call: call, - Workdir: activeWorkdir, - ToolTimeout: time.Duration(cfg.ToolTimeoutSec) * time.Second, - }) - if s.isRunCanceled(execErr) { - return s.handleRunError(ctx, input.RunID, session.ID, execErr) - } - if execErr == nil { - if err := ctx.Err(); err != nil { - return s.handleRunError(ctx, input.RunID, session.ID, err) - } - } - - if execErr != nil && strings.TrimSpace(result.Content) == "" { - result.Content = execErr.Error() - } - if permissionEvent, ok := permissionEventFromError(execErr); ok { - s.emit(ctx, EventPermissionResolved, input.RunID, session.ID, permissionEvent.toResolvedPayload()) - } - - toolMessage := providertypes.Message{ - Role: providertypes.RoleTool, - Content: result.Content, - ToolCallID: call.ID, - IsError: result.IsError, - ToolMetadata: tools.SanitizeToolMetadata(result.Name, result.Metadata), - } - session.Messages = append(session.Messages, toolMessage) - session.UpdatedAt = time.Now() - if err := s.sessionStore.Save(ctx, &session); err != nil { - if execErr != nil && errors.Is(err, context.Canceled) { - s.emit(ctx, EventToolResult, input.RunID, session.ID, result) - } - return s.handleRunError(ctx, input.RunID, session.ID, err) - } - if err := ctx.Err(); err != nil { - if execErr == nil { - return s.handleRunError(ctx, input.RunID, session.ID, err) - } - } - - s.emit(ctx, EventToolResult, input.RunID, session.ID, result) - if execErr != nil { - if err := ctx.Err(); err != nil { - return s.handleRunError(ctx, input.RunID, session.ID, err) - } - } - } - } -} - +// CancelActiveRun 尝试取消当前正在执行的 Run。 func (s *Service) CancelActiveRun() bool { s.runMu.Lock() cancel := s.activeRunCancel @@ -405,14 +109,17 @@ func (s *Service) CancelActiveRun() bool { return true } +// Events 返回 runtime 事件通道，供上层 UI 订阅。 func (s *Service) Events() <-chan RuntimeEvent { return s.events } +// ListSessions 返回当前会话存储中的所有摘要。 func (s *Service) ListSessions(ctx context.Context) ([]agentsession.Summary, error) { return s.sessionStore.ListSummaries(ctx) } +// LoadSession 按 id 加载完整会话内容。 func (s *Service) LoadSession(ctx context.Context, id string) (agentsession.Session, error) { session, err := s.sessionStore.Load(ctx, id) if err != nil { @@ -421,6 +128,7 @@ func (s *Service) LoadSession(ctx context.Context, id string) (agentsession.Sess return session, nil } +// loadOrCreateSession 负责在运行开始时解析工作目录并加载或创建会话。 func (s *Service) loadOrCreateSession( ctx context.Context, sessionID string, @@ -439,6 +147,7 @@ func (s *Service) loadOrCreateSession( } return session, nil } + session, err := s.sessionStore.Load(ctx, sessionID) if err != nil { return agentsession.Session{}, err @@ -454,6 +163,7 @@ func (s *Service) loadOrCreateSession( if session.Workdir == resolved { return session, nil } + session.Workdir = resolved session.UpdatedAt = time.Now() if err := s.sessionStore.Save(ctx, &session); err != nil { @@ -462,9 +172,7 @@ func (s *Service) loadOrCreateSession( return session, nil } -// emit 向事件通道发送事件，并在通道阻塞且上下文取消时返回对应错误。 -// 先尝试非阻塞发送，确保即使 context 已取消，只要通道有空间事件就能被投递； -// 仅在通道已满时才通过 ctx.Done() 退出，避免 goroutine 泄漏并向调用方反馈未投递状态。 +// emit 将 runtime 事件投递到事件通道，并在通道阻塞且上下文取消时返回错误。 func (s *Service) emit(ctx context.Context, kind EventType, runID string, sessionID string, payload any) error { evt := RuntimeEvent{ Type: kind, @@ -485,131 +193,7 @@ func (s *Service) emit(ctx context.Context, kind EventType, runID string, sessio } } -// handleProviderStreamEvent 解析并应用单条 provider 流式事件，缺失载荷或未知类型时返回错误。 -func handleProviderStreamEvent( - event providertypes.StreamEvent, - acc *streamAccumulator, - onTextDelta func(string), - onToolCallStart func(providertypes.ToolCallStartPayload), - onMessageDone func(providertypes.MessageDonePayload), -) error { - switch event.Type { - case providertypes.StreamEventTextDelta: - payload, err := event.TextDeltaValue() - if err != nil { - return err - } - if onTextDelta != nil { - onTextDelta(payload.Text) - } - if acc != nil { - acc.accumulateTextDelta(payload.Text) - } - case providertypes.StreamEventToolCallStart: - payload, err := event.ToolCallStartValue() - if err != nil { - return err - } - if onToolCallStart != nil { - onToolCallStart(payload) - } - if acc != nil { - acc.accumulateToolCallStart(payload.Index, payload.ID, payload.Name) - } - case providertypes.StreamEventToolCallDelta: - payload, err := event.ToolCallDeltaValue() - if err != nil { - return err - } - if acc != nil { - acc.accumulateToolCallDelta(payload.Index, payload.ID, payload.ArgumentsDelta) - } - case providertypes.StreamEventMessageDone: - payload, err := event.MessageDoneValue() - if err != nil { - return err - } - if acc != nil { - acc.messageDone = true - } - if onMessageDone != nil { - onMessageDone(payload) - } - default: - return fmt.Errorf("runtime: unsupported provider stream event type %q", event.Type) - } - return nil -} - -// forwardProviderEvents 将 provider 流式事件转发为 runtime 事件，同时向 accumulator 累积消息状态。 -// 使用 select 同时监听输入通道和 context 取消信号，确保 goroutine 不会因通道阻塞而泄漏。 -func (s *Service) forwardProviderEvents( - ctx context.Context, - runID string, - sessionID string, - input <-chan providertypes.StreamEvent, - done chan<- error, - acc *streamAccumulator, -) { - var forwardErr error - defer func() { - done <- forwardErr - }() - - for { - select { - case event, ok := <-input: - if !ok { - return - } - err := handleProviderStreamEvent( - event, - acc, - func(text string) { - s.emit(ctx, EventAgentChunk, runID, sessionID, text) - }, - func(payload providertypes.ToolCallStartPayload) { - s.emit(ctx, EventToolCallThinking, runID, sessionID, payload.Name) - }, - func(done providertypes.MessageDonePayload) { - if done.Usage != nil { - s.sessionInputTokens += done.Usage.InputTokens - s.sessionOutputTokens += done.Usage.OutputTokens - s.emit(ctx, EventTokenUsage, runID, sessionID, TokenUsagePayload{ - InputTokens: done.Usage.InputTokens, - OutputTokens: done.Usage.OutputTokens, - SessionInputTokens: s.sessionInputTokens, - SessionOutputTokens: s.sessionOutputTokens, - }) - } - }, - ) - if err != nil && forwardErr == nil { - // 记录首个协议错误后继续排空事件通道，避免 provider 在后续发送时阻塞。 - forwardErr = err - } - case <-ctx.Done(): - return - } - } -} - -// restoreSessionTokens restores token counters from session persistence, -// or resets them to zero for new sessions. -func (s *Service) restoreSessionTokens(session agentsession.Session) { - s.sessionInputTokens = session.TokenInputTotal - s.sessionOutputTokens = session.TokenOutputTotal -} - -// autoCompactThreshold returns the configured auto-compact input token threshold, -// or 0 if auto-compact is disabled. -func (s *Service) autoCompactThreshold(cfg config.Config) int { - if cfg.Context.AutoCompact.Enabled && cfg.Context.AutoCompact.InputTokenThreshold > 0 { - return cfg.Context.AutoCompact.InputTokenThreshold - } - return 0 -} - +// startRun 记录当前激活的运行取消句柄，并分配一个新的运行令牌。 func (s *Service) startRun(cancel context.CancelFunc) uint64 { s.runMu.Lock() defer s.runMu.Unlock() @@ -621,6 +205,7 @@ func (s *Service) startRun(cancel context.CancelFunc) uint64 { return token } +// finishRun 在运行结束时释放当前激活的取消句柄。 func (s *Service) finishRun(token uint64) { s.runMu.Lock() defer s.runMu.Unlock() @@ -633,214 +218,22 @@ func (s *Service) finishRun(token uint64) { s.activeRunCancel = nil } -func (s *Service) handleRunError(ctx context.Context, runID string, sessionID string, err error) error { - if s.isRunCanceled(err) { - s.emit(ctx, EventRunCanceled, runID, sessionID, nil) - return context.Canceled - } - - // 提取 ProviderError 信息用于日志增强。 - var pErr *provider.ProviderError - if errors.As(err, &pErr) { - log.Printf("runtime: provider error (status=%d, code=%s, retryable=%v): %s", - pErr.StatusCode, pErr.Code, pErr.Retryable, pErr.Message) - } - - s.emit(ctx, EventError, runID, sessionID, err.Error()) - return err -} - -// isRetryableProviderError 判断 error 是否为可重试的 ProviderError。 -func isRetryableProviderError(err error) bool { - var pErr *provider.ProviderError - if !errors.As(err, &pErr) { - return false - } - return pErr.Retryable -} - -// callProviderWithRetry 在可重试的 ProviderError 上自动重试 provider.Generate() 调用。 -// 每次重试都会重新创建 provider 实例、流式事件转发管道和累积器。 -// 非可重试错误、context 取消、重试耗尽时直接返回错误。 -// 返回值 acc 保存了本轮流式事件的完整累积状态，供调用方构建 assistant Message 使用。 -func (s *Service) callProviderWithRetry( - ctx context.Context, - runID string, - sessionID string, - req providertypes.GenerateRequest, -) (*streamAccumulator, error) { - acc := newStreamAccumulator() - var lastErr error - - for retryAttempt := 0; retryAttempt <= defaultProviderRetryMax; retryAttempt++ { - if retryAttempt > 0 { - // 重试时重置累积器，避免混入上轮数据 - acc = newStreamAccumulator() - wait := providerRetryBackoff(retryAttempt) - s.emit(ctx, EventProviderRetry, runID, sessionID, - fmt.Sprintf("retrying provider call (attempt %d/%d, wait=%.1fs)...", - retryAttempt, defaultProviderRetryMax, wait.Seconds())) - - select { - case <-ctx.Done(): - return nil, ctx.Err() - case <-time.After(wait): - } - } - - cfg := s.configManager.Get() - resolvedProvider, err := config.ResolveSelectedProvider(cfg) - if err != nil { - return nil, err - } - runtimeCfg := resolvedProvider.ToRuntimeConfig() - - modelProvider, err := s.providerFactory.Build(ctx, runtimeCfg) - if err != nil { - return nil, err - } - - streamEvents := make(chan providertypes.StreamEvent, 32) - streamDone := make(chan error, 1) - go s.forwardProviderEvents(ctx, runID, sessionID, streamEvents, streamDone, acc) - - err = modelProvider.Generate(ctx, req, streamEvents) - close(streamEvents) - forwardErr := <-streamDone - if forwardErr != nil { - if err != nil { - return nil, fmt.Errorf("runtime: provider stream handling failed after provider error: %v: %w", err, forwardErr) - } - return nil, forwardErr - } - - if err == nil { - if !acc.messageDone { - return nil, fmt.Errorf("%w: provider stream ended without message_done event", provider.ErrStreamInterrupted) - } - return acc, nil - } - - lastErr = err - - // 非可重试错误或 context 已取消，立即返回。 - if !isRetryableProviderError(err) { - return nil, lastErr - } - if ctx.Err() != nil { - return nil, ctx.Err() - } - } - - if lastErr == nil { - lastErr = errors.New("max retries exceeded") - } - return nil, fmt.Errorf("runtime: max retries exhausted, last error: %w", lastErr) -} - -// providerRetryBackoff 计算指数退避 + 随机抖动的等待时间。 -// attempt 从 1 开始（首次重试）。 -func providerRetryBackoff(attempt int) time.Duration { - wait := providerRetryBaseWait << (attempt - 1) - - // 随机抖动：[0.5, 1.5) * wait - jitter := float64(wait) * (0.5 + rand.Float64()) - wait = time.Duration(jitter) - - if wait > providerRetryMaxWait { - wait = providerRetryMaxWait - } - - return wait -} - -func (s *Service) isRunCanceled(err error) bool { - return errors.Is(err, context.Canceled) -} - -type permissionEventView struct { - toolName string - actionType string - operation string - targetType string - target string - decision string - reason string - ruleID string - scope string - resolvedAs string -} +// acquireSessionLock 获取指定会话的互斥锁，确保同一会话的 Run/Compact 串行执行。 +// 不同会话的锁互不干扰，支持多会话并行。 +func (s *Service) acquireSessionLock(sessionID string) *sync.Mutex { + s.sessionMu.Lock() + defer s.sessionMu.Unlock() -// permissionEventFromError 从工具权限错误中提取可用于 runtime 事件的结构化信息。 -func permissionEventFromError(err error) (permissionEventView, bool) { - var permissionErr *tools.PermissionDecisionError - if !errors.As(err, &permissionErr) { - return permissionEventView{}, false + mu, ok := s.sessionLocks[sessionID] + if !ok { + mu = &sync.Mutex{} + s.sessionLocks[sessionID] = mu } - - action := permissionErr.Action() - decision := strings.TrimSpace(permissionErr.Decision()) - reason := strings.TrimSpace(permissionErr.Reason()) - if reason == "" { - reason = strings.TrimSpace(err.Error()) - } - - resolvedAs := "denied" - if strings.EqualFold(decision, "ask") { - resolvedAs = "rejected" - } - - return permissionEventView{ - toolName: action.Payload.ToolName, - actionType: string(action.Type), - operation: action.Payload.Operation, - targetType: string(action.Payload.TargetType), - target: action.Payload.Target, - decision: decision, - reason: reason, - ruleID: strings.TrimSpace(permissionErr.RuleID()), - scope: strings.TrimSpace(permissionErr.RememberScope()), - resolvedAs: resolvedAs, - }, true -} - -// toRequestPayload 将权限事件视图转换为 permission_request 载荷。 -func (v permissionEventView) toRequestPayload() PermissionRequestPayload { - return PermissionRequestPayload{ - ToolName: v.toolName, - ActionType: v.actionType, - Operation: v.operation, - TargetType: v.targetType, - Target: v.target, - Decision: v.decision, - Reason: v.reason, - RuleID: v.ruleID, - RememberScope: v.scope, - } -} - -// toResolvedPayload 将权限事件视图转换为 permission_resolved 载荷。 -func (v permissionEventView) toResolvedPayload() PermissionResolvedPayload { - return PermissionResolvedPayload{ - ToolName: v.toolName, - ActionType: v.actionType, - Operation: v.operation, - TargetType: v.targetType, - Target: v.target, - Decision: v.decision, - Reason: v.reason, - RuleID: v.ruleID, - RememberScope: v.scope, - ResolvedAs: v.resolvedAs, - } -} - -func effectiveSessionWorkdir(sessionWorkdir string, defaultWorkdir string) string { - return agentsession.EffectiveWorkdir(sessionWorkdir, defaultWorkdir) + return mu } func resolveWorkdirForSession(defaultWorkdir string, currentWorkdir string, requestedWorkdir string) (string, error) { - base := effectiveSessionWorkdir(currentWorkdir, defaultWorkdir) + base := agentsession.EffectiveWorkdir(currentWorkdir, defaultWorkdir) if strings.TrimSpace(requestedWorkdir) == "" { return agentsession.ResolveExistingDir(base) } @@ -851,7 +244,3 @@ func resolveWorkdirForSession(defaultWorkdir string, currentWorkdir string, requ } return agentsession.ResolveExistingDir(target) } - -func normalizeExistingWorkdir(workdir string) (string, error) { - return agentsession.ResolveExistingDir(workdir) -} diff --git a/internal/runtime/runtime_gap_coverage_test.go b/internal/runtime/runtime_gap_coverage_test.go new file mode 100644 index 00000000..6eb03dcc --- /dev/null +++ b/internal/runtime/runtime_gap_coverage_test.go @@ -0,0 +1,288 @@ +package runtime + +import ( + "context" + "errors" + "sync" + "testing" + "time" + + "neo-code/internal/config" + contextcompact "neo-code/internal/context/compact" + "neo-code/internal/provider" + providertypes "neo-code/internal/provider/types" + approvalflow "neo-code/internal/runtime/approval" + "neo-code/internal/security" + agentsession "neo-code/internal/session" + "neo-code/internal/tools" +) + +func TestCompactValidationAndLoadErrorBranches(t *testing.T) { + t.Parallel() + + service := &Service{ + configManager: newRuntimeConfigManager(t), + sessionStore: newMemoryStore(), + events: make(chan RuntimeEvent, 16), + sessionLocks: make(map[string]*sync.Mutex), + } + + ctx, cancel := context.WithCancel(context.Background()) + cancel() + if _, err := service.Compact(ctx, CompactInput{SessionID: "s-1"}); !errors.Is(err, context.Canceled) { + t.Fatalf("expected context canceled, got %v", err) + } + + if _, err := service.Compact(context.Background(), CompactInput{}); err == nil { + t.Fatalf("expected empty session id error") + } + + if _, err := service.Compact(context.Background(), CompactInput{SessionID: "missing"}); err == nil { + t.Fatalf("expected load error for missing session") + } +} + +func TestRunCompactForSessionPolicyBranches(t *testing.T) { + t.Parallel() + + session := newRuntimeSession("session-compact-policy") + session.Messages = []providertypes.Message{{Role: providertypes.RoleUser, Content: "before"}} + + service := &Service{ + events: make(chan RuntimeEvent, 32), + } + + // default runner init failure: strict returns error + if _, _, err := service.runCompactForSession(context.Background(), "run-compact", session, config.Config{}, contextcompact.ModeManual, compactErrorStrict); err == nil { + t.Fatalf("expected strict mode to return default runner error") + } + + // default runner init failure: best effort swallows error + if gotSession, gotResult, err := service.runCompactForSession(context.Background(), "run-compact", session, config.Config{}, contextcompact.ModeManual, compactErrorBestEffort); err != nil { + t.Fatalf("expected best effort to swallow runner init error, got %v", err) + } else if gotResult.Applied || len(gotSession.Messages) != len(session.Messages) { + t.Fatalf("expected noop result in best effort mode") + } + + service.compactRunner = &stubCompactRunner{err: errors.New("runner failed")} + if _, _, err := service.runCompactForSession(context.Background(), "run-compact", session, config.Config{}, contextcompact.ModeManual, compactErrorStrict); err == nil { + t.Fatalf("expected strict mode to return runner error") + } + if _, _, err := service.runCompactForSession(context.Background(), "run-compact", session, config.Config{}, contextcompact.ModeManual, compactErrorBestEffort); err != nil { + t.Fatalf("expected best effort to swallow runner error, got %v", err) + } +} + +func TestRunCompactForSessionSaveErrorPolicyBranches(t *testing.T) { + t.Parallel() + + baseStore := newMemoryStore() + session := newRuntimeSession("session-compact-save-error") + session.Messages = []providertypes.Message{{Role: providertypes.RoleUser, Content: "before"}} + baseStore.sessions[session.ID] = cloneSession(session) + + store := &failingStore{Store: baseStore, saveErr: errors.New("save failed"), failOnSave: 1, ignoreContextErr: true} + service := &Service{ + sessionStore: store, + events: make(chan RuntimeEvent, 16), + compactRunner: &stubCompactRunner{result: contextcompact.Result{ + Applied: true, + Messages: []providertypes.Message{{Role: providertypes.RoleAssistant, Content: "after"}}, + }}, + } + + strictSession, _, err := service.runCompactForSession(context.Background(), "run-compact-save", session, config.Config{}, contextcompact.ModeManual, compactErrorStrict) + if err == nil { + t.Fatalf("expected strict mode save error") + } + if strictSession.Messages[0].Content != "before" { + t.Fatalf("expected strict mode to rollback messages, got %+v", strictSession.Messages) + } + + store.saveCalls = 0 + bestEffortSession, bestEffortResult, err := service.runCompactForSession(context.Background(), "run-compact-save", session, config.Config{}, contextcompact.ModeManual, compactErrorBestEffort) + if err != nil { + t.Fatalf("expected best effort to swallow save error, got %v", err) + } + if bestEffortResult.Applied { + t.Fatalf("expected empty compact result on best effort save failure") + } + if bestEffortSession.Messages[0].Content != "before" { + t.Fatalf("expected best effort rollback messages, got %+v", bestEffortSession.Messages) + } +} + +func TestCompactProviderSelectionErrorBranches(t *testing.T) { + t.Parallel() + + manager := newRuntimeConfigManager(t) + cfg := manager.Get() + session := newRuntimeSession("session-provider-select") + session.Provider = "missing-provider" + session.Model = "m1" + if _, _, err := resolveCompactProviderSelection(session, cfg); err == nil { + t.Fatalf("expected provider not found error") + } + + cfg.SelectedProvider = "" + if _, _, err := resolveCompactProviderSelection(agentsession.Session{}, cfg); err == nil { + t.Fatalf("expected selected provider empty error") + } + + service := &Service{providerFactory: &scriptedProviderFactory{provider: &scriptedProvider{}}} + if _, err := service.defaultCompactRunner(agentsession.Session{}, config.Config{}); err == nil { + t.Fatalf("expected defaultCompactRunner to fail on invalid config") + } +} + +func TestCompactSummaryGeneratorErrorBranches(t *testing.T) { + t.Parallel() + + ctx, cancel := context.WithCancel(context.Background()) + cancel() + g := &compactSummaryGenerator{} + if _, err := g.Generate(ctx, contextcompact.SummaryInput{}); !errors.Is(err, context.Canceled) { + t.Fatalf("expected canceled error, got %v", err) + } + + if _, err := (&compactSummaryGenerator{}).Generate(context.Background(), contextcompact.SummaryInput{}); err == nil { + t.Fatalf("expected nil provider factory error") + } + + if _, err := (&compactSummaryGenerator{providerFactory: &scriptedProviderFactory{provider: &scriptedProvider{}}}).Generate(context.Background(), contextcompact.SummaryInput{}); err == nil { + t.Fatalf("expected incomplete provider config error") + } + + g = &compactSummaryGenerator{ + providerFactory: &scriptedProviderFactory{err: errors.New("build failed")}, + providerConfig: provider.RuntimeConfig{Name: "openai", Driver: "openai", BaseURL: "https://example.com", APIKey: "k"}, + } + if _, err := g.Generate(context.Background(), contextcompact.SummaryInput{}); err == nil { + t.Fatalf("expected provider build error") + } + + g = &compactSummaryGenerator{ + providerFactory: &scriptedProviderFactory{provider: &scriptedProvider{streams: [][]providertypes.StreamEvent{{providertypes.NewTextDeltaStreamEvent(" ")}}}}, + providerConfig: provider.RuntimeConfig{Name: "openai", Driver: "openai", BaseURL: "https://example.com", APIKey: "k"}, + } + if _, err := g.Generate(context.Background(), contextcompact.SummaryInput{}); err == nil { + t.Fatalf("expected empty summary error") + } +} + +func TestRuntimeLoadOrCreateAndEmitBranches(t *testing.T) { + t.Parallel() + + manager := newRuntimeConfigManager(t) + service := &Service{configManager: manager, sessionStore: newMemoryStore(), events: make(chan RuntimeEvent, 1), sessionLocks: map[string]*sync.Mutex{}} + + if _, err := service.loadOrCreateSession(context.Background(), "", "title", t.TempDir(), "missing/subdir"); err == nil { + t.Fatalf("expected new session workdir resolve error") + } + + if _, err := service.loadOrCreateSession(context.Background(), "not-found", "title", t.TempDir(), ""); err == nil { + t.Fatalf("expected load existing session error") + } + + store := newMemoryStore() + session := newRuntimeSession("session-load-update") + session.Workdir = t.TempDir() + store.sessions[session.ID] = cloneSession(session) + service.sessionStore = store + if _, err := service.loadOrCreateSession(context.Background(), session.ID, "title", t.TempDir(), "missing/child"); err == nil { + t.Fatalf("expected resolve error for requested workdir") + } + + store.saves = 0 + if got, err := service.loadOrCreateSession(context.Background(), session.ID, "title", t.TempDir(), "."); err != nil { + t.Fatalf("loadOrCreateSession same-workdir error = %v", err) + } else if got.Workdir != session.Workdir || store.saves != 0 { + t.Fatalf("expected same workdir path to skip save, got workdir=%q saves=%d", got.Workdir, store.saves) + } + + fStore := &failingStore{Store: store, saveErr: errors.New("save failed"), failOnSave: 1, ignoreContextErr: true} + service.sessionStore = fStore + if _, err := service.loadOrCreateSession(context.Background(), session.ID, "title", t.TempDir(), ".."); err == nil { + t.Fatalf("expected save error when updating session workdir") + } + + service.events <- RuntimeEvent{Type: EventAgentChunk} + done := make(chan struct{}) + go func() { + <-time.After(10 * time.Millisecond) + <-service.events + close(done) + }() + if err := service.emit(context.Background(), EventAgentDone, "run", "session", "ok"); err != nil { + t.Fatalf("emit() error = %v", err) + } + <-done + + token1 := service.startRun(func() {}) + token2 := service.startRun(func() {}) + service.finishRun(token1) + if service.activeRunToken != token2 { + t.Fatalf("expected finishRun with stale token to keep active run") + } +} + +func TestPermissionHelperAndAwaitBranches(t *testing.T) { + t.Parallel() + + if status := permissionResolutionStatus(string(security.DecisionAsk)); status != permissionResolvedRejected { + t.Fatalf("expected ask to map rejected, got %q", status) + } + if category := permissionToolCategory(security.Action{Type: security.ActionTypeWrite, Payload: security.ActionPayload{Resource: "filesystem_write"}}); category != permissionToolCategoryFilesystemWrite { + t.Fatalf("expected filesystem write category, got %q", category) + } + if category := permissionToolCategory(security.Action{Payload: security.ActionPayload{ToolName: "fallback_tool"}}); category != "fallback_tool" { + t.Fatalf("expected tool_name fallback, got %q", category) + } + + service := &Service{approvalBroker: nil, events: make(chan RuntimeEvent, 8)} + permissionErr := permissionDecisionAskError(t) + _, _, err := service.awaitPermissionDecision(context.Background(), permissionExecutionInput{RunID: "r", SessionID: "s", Call: providertypes.ToolCall{ID: "c", Name: "filesystem_read_file"}}, permissionErr) + if err == nil { + t.Fatalf("expected open error from nil approval broker") + } + + service.approvalBroker = approvalflow.NewBroker() + ctx, cancel := context.WithCancel(context.Background()) + cancel() + _, _, err = service.awaitPermissionDecision(ctx, permissionExecutionInput{RunID: "r", SessionID: "s", Call: providertypes.ToolCall{ID: "c", Name: "filesystem_read_file"}}, permissionErr) + if !errors.Is(err, context.Canceled) { + t.Fatalf("expected canceled while waiting permission, got %v", err) + } +} + +func permissionDecisionAskError(t *testing.T) *tools.PermissionDecisionError { + t.Helper() + + registry := tools.NewRegistry() + registry.Register(&stubTool{name: "filesystem_read_file", content: "ok"}) + engine, err := security.NewStaticGateway(security.DecisionAllow, []security.Rule{{ + ID: "ask-read", + Type: security.ActionTypeRead, + Resource: "filesystem_read_file", + Decision: security.DecisionAsk, + }}) + if err != nil { + t.Fatalf("new static gateway: %v", err) + } + manager, err := tools.NewManager(registry, engine, nil) + if err != nil { + t.Fatalf("new tool manager: %v", err) + } + _, execErr := manager.Execute(context.Background(), tools.ToolCallInput{ + ID: "call-ask", + Name: "filesystem_read_file", + Arguments: []byte(`{"path":"README.md"}`), + Workdir: t.TempDir(), + SessionID: "session-ask", + }) + var permissionErr *tools.PermissionDecisionError + if !errors.As(execErr, &permissionErr) { + t.Fatalf("expected PermissionDecisionError, got %v", execErr) + } + return permissionErr +} diff --git a/internal/runtime/runtime_internal_helpers_test.go b/internal/runtime/runtime_internal_helpers_test.go new file mode 100644 index 00000000..c9d72647 --- /dev/null +++ b/internal/runtime/runtime_internal_helpers_test.go @@ -0,0 +1,241 @@ +package runtime + +import ( + "context" + "errors" + "testing" + "time" + + "neo-code/internal/config" + "neo-code/internal/provider" + providertypes "neo-code/internal/provider/types" + "neo-code/internal/runtime/approval" + agentsession "neo-code/internal/session" + "neo-code/internal/tools" +) + +func TestRunStateNilReceiverNoops(t *testing.T) { + t.Parallel() + + var state *runState + state.syncSessionTokenTotals() + state.recordUsage(3, 5) + state.resetTokenTotals() + state.touchSession() +} + +func TestRunStateMutationsAndSync(t *testing.T) { + t.Parallel() + + session := newRuntimeSession("session-state") + state := newRunState("run-1", session) + + state.recordUsage(10, 20) + if state.tokenInputTotal != 11 || state.tokenOutputTotal != 22 { + t.Fatalf("unexpected token totals: in=%d out=%d", state.tokenInputTotal, state.tokenOutputTotal) + } + + state.syncSessionTokenTotals() + if state.session.TokenInputTotal != 11 || state.session.TokenOutputTotal != 22 { + t.Fatalf("session totals not synced: %+v", state.session) + } + + state.resetTokenTotals() + if state.tokenInputTotal != 0 || state.tokenOutputTotal != 0 { + t.Fatalf("expected reset totals to be zero, got in=%d out=%d", state.tokenInputTotal, state.tokenOutputTotal) + } + + before := state.session.UpdatedAt + state.recordUsage(1, 2) + state.touchSession() + if state.session.UpdatedAt.Before(before) { + t.Fatalf("expected touchSession to update time") + } + if state.session.TokenInputTotal != 1 || state.session.TokenOutputTotal != 2 { + t.Fatalf("expected touchSession to sync totals") + } +} + +func TestAppendAssistantMessageAndSaveMetadataBranches(t *testing.T) { + t.Parallel() + + store := newMemoryStore() + session := newRuntimeSession("session-append-assistant") + store.sessions[session.ID] = cloneSession(session) + + service := &Service{sessionStore: store} + state := newRunState("run-append-assistant", session) + snapshot := turnSnapshot{ + providerConfig: providerRuntimeConfigForTest("openai"), + model: "gpt-4.1", + } + + if err := service.appendAssistantMessageAndSave(context.Background(), &state, snapshot, providertypes.Message{Role: providertypes.RoleAssistant}); err != nil { + t.Fatalf("appendAssistantMessageAndSave() error = %v", err) + } + if store.saves != 1 { + t.Fatalf("expected metadata change to persist once, saves=%d", store.saves) + } + + store.saves = 0 + state.session.Provider = snapshot.providerConfig.Name + state.session.Model = snapshot.model + if err := service.appendAssistantMessageAndSave(context.Background(), &state, snapshot, providertypes.Message{Role: providertypes.RoleAssistant}); err != nil { + t.Fatalf("appendAssistantMessageAndSave() error = %v", err) + } + if store.saves != 0 { + t.Fatalf("expected empty assistant without metadata change to skip save, saves=%d", store.saves) + } +} + +func TestAppendToolMessageAndSaveSanitizesMetadata(t *testing.T) { + t.Parallel() + + store := newMemoryStore() + session := newRuntimeSession("session-append-tool") + store.sessions[session.ID] = cloneSession(session) + + service := &Service{sessionStore: store} + state := newRunState("run-append-tool", session) + call := providertypes.ToolCall{ID: "call-1", Name: "filesystem_read_file"} + result := tools.ToolResult{ + Name: "filesystem_read_file", + Content: "ok", + Metadata: map[string]any{ + "stderr": "warn", + "sensitive": "drop-me", + }, + } + if err := service.appendToolMessageAndSave(context.Background(), &state, call, result); err != nil { + t.Fatalf("appendToolMessageAndSave() error = %v", err) + } + if len(state.session.Messages) != 1 { + t.Fatalf("expected one persisted tool message, got %d", len(state.session.Messages)) + } + msg := state.session.Messages[0] + if msg.ToolMetadata["tool_name"] != "filesystem_read_file" { + t.Fatalf("expected tool_name metadata, got %+v", msg.ToolMetadata) + } + if _, exists := msg.ToolMetadata["sensitive"]; exists { + t.Fatalf("expected sensitive metadata key to be removed, got %+v", msg.ToolMetadata) + } +} + +func TestResolveMaxLoopsBranches(t *testing.T) { + t.Parallel() + + if got := resolveMaxLoops(config.Config{MaxLoops: 0}); got != defaultMaxLoops { + t.Fatalf("expected default max loops for zero, got %d", got) + } + if got := resolveMaxLoops(config.Config{MaxLoops: -3}); got != defaultMaxLoops { + t.Fatalf("expected default max loops for negative, got %d", got) + } + if got := resolveMaxLoops(config.Config{MaxLoops: 12}); got != 12 { + t.Fatalf("expected explicit max loops, got %d", got) + } +} + +func TestEmitTokenUsageSkipsZeroUsage(t *testing.T) { + t.Parallel() + + service := &Service{events: make(chan RuntimeEvent, 8)} + state := &runState{runID: "run-token", session: newRuntimeSession("session-token")} + + service.emitTokenUsage(context.Background(), state, providerTurnResult{}) + events := collectRuntimeEvents(service.Events()) + if len(events) != 0 { + t.Fatalf("expected no token event for zero usage, got %+v", events) + } + + state.recordUsage(5, 7) + service.emitTokenUsage(context.Background(), state, providerTurnResult{inputTokens: 5, outputTokens: 7}) + events = collectRuntimeEvents(service.Events()) + if len(events) != 1 || events[0].Type != EventTokenUsage { + t.Fatalf("expected one token usage event, got %+v", events) + } +} + +func TestExecuteAssistantToolCallsFillsErrorContent(t *testing.T) { + t.Parallel() + + store := newMemoryStore() + session := newRuntimeSession("session-exec-tool-error-fill") + store.sessions[session.ID] = cloneSession(session) + + toolErr := errors.New("tool exploded") + manager := &stubToolManager{err: toolErr} + service := &Service{ + sessionStore: store, + toolManager: manager, + approvalBroker: approval.NewBroker(), + events: make(chan RuntimeEvent, 32), + } + state := newRunState("run-exec-tool-error-fill", session) + assistant := providertypes.Message{ + Role: providertypes.RoleAssistant, + ToolCalls: []providertypes.ToolCall{ + {ID: "call-err", Name: "filesystem_read_file", Arguments: `{"path":"a.txt"}`}, + }, + } + snapshot := turnSnapshot{workdir: t.TempDir(), toolTimeout: time.Second} + + if err := service.executeAssistantToolCalls(context.Background(), &state, snapshot, assistant); err != nil { + t.Fatalf("executeAssistantToolCalls() error = %v", err) + } + if len(state.session.Messages) != 1 { + t.Fatalf("expected one tool message, got %d", len(state.session.Messages)) + } + if state.session.Messages[0].Content != toolErr.Error() { + t.Fatalf("expected tool error content fallback, got %q", state.session.Messages[0].Content) + } +} + +func TestExecuteAssistantToolCallsCanceledSaveStillEmitsResultWhenExecErr(t *testing.T) { + t.Parallel() + + baseStore := newMemoryStore() + session := newRuntimeSession("session-exec-tool-cancel-save") + baseStore.sessions[session.ID] = cloneSession(session) + store := &failingStore{ + Store: baseStore, + saveErr: context.Canceled, + failOnSave: 1, + ignoreContextErr: true, + } + + manager := &stubToolManager{err: errors.New("tool failed")} + service := &Service{ + sessionStore: store, + toolManager: manager, + approvalBroker: approval.NewBroker(), + events: make(chan RuntimeEvent, 32), + } + state := newRunState("run-exec-tool-cancel-save", session) + assistant := providertypes.Message{ + Role: providertypes.RoleAssistant, + ToolCalls: []providertypes.ToolCall{ + {ID: "call-1", Name: "filesystem_read_file", Arguments: `{"path":"a.txt"}`}, + }, + } + snapshot := turnSnapshot{workdir: t.TempDir(), toolTimeout: time.Second} + + err := service.executeAssistantToolCalls(context.Background(), &state, snapshot, assistant) + if !errors.Is(err, context.Canceled) { + t.Fatalf("expected context.Canceled from save failure, got %v", err) + } + + events := collectRuntimeEvents(service.Events()) + assertEventSequence(t, events, []EventType{EventToolStart, EventToolResult}) +} + +func newRuntimeSession(id string) agentsession.Session { + session := agentsession.New("runtime test") + session.ID = id + session.TokenInputTotal = 1 + session.TokenOutputTotal = 2 + return session +} + +func providerRuntimeConfigForTest(name string) provider.RuntimeConfig { + return provider.RuntimeConfig{Name: name} +} diff --git a/internal/runtime/runtime_remaining_branches_test.go b/internal/runtime/runtime_remaining_branches_test.go new file mode 100644 index 00000000..d618a5f1 --- /dev/null +++ b/internal/runtime/runtime_remaining_branches_test.go @@ -0,0 +1,478 @@ +package runtime + +import ( + "context" + "errors" + "os" + "sync" + "testing" + "time" + + "neo-code/internal/config" + agentcontext "neo-code/internal/context" + contextcompact "neo-code/internal/context/compact" + "neo-code/internal/provider" + "neo-code/internal/provider/streaming" + + providertypes "neo-code/internal/provider/types" + approvalflow "neo-code/internal/runtime/approval" + "neo-code/internal/security" + agentsession "neo-code/internal/session" + "neo-code/internal/tools" +) + +type callbackToolManager struct { + executeFn func(ctx context.Context, input tools.ToolCallInput) (tools.ToolResult, error) + rememberErr error +} + +func (m *callbackToolManager) ListAvailableSpecs(ctx context.Context, input tools.SpecListInput) ([]providertypes.ToolSpec, error) { + return nil, ctx.Err() +} + +func (m *callbackToolManager) MicroCompactPolicy(name string) tools.MicroCompactPolicy { + return tools.MicroCompactPolicyCompact +} + +func (m *callbackToolManager) Execute(ctx context.Context, input tools.ToolCallInput) (tools.ToolResult, error) { + if m.executeFn != nil { + return m.executeFn(ctx, input) + } + return tools.ToolResult{Name: input.Name}, nil +} + +func (m *callbackToolManager) RememberSessionDecision(sessionID string, action security.Action, scope tools.SessionPermissionScope) error { + return m.rememberErr +} + +type saveHookStore struct { + base *memoryStore + saveHook func() +} + +func (s *saveHookStore) Save(ctx context.Context, session *agentsession.Session) error { + if s.saveHook != nil { + s.saveHook() + } + return s.base.Save(ctx, session) +} + +type postSaveHookStore struct { + base *memoryStore + saveHook func() + once sync.Once +} + +func (s *postSaveHookStore) Save(ctx context.Context, session *agentsession.Session) error { + err := s.base.Save(ctx, session) + if err == nil && s.saveHook != nil { + s.once.Do(s.saveHook) + } + return err +} + +func (s *saveHookStore) Load(ctx context.Context, id string) (agentsession.Session, error) { + return s.base.Load(ctx, id) +} + +func (s *saveHookStore) ListSummaries(ctx context.Context) ([]agentsession.Summary, error) { + return s.base.ListSummaries(ctx) +} + +func (s *postSaveHookStore) Load(ctx context.Context, id string) (agentsession.Session, error) { + return s.base.Load(ctx, id) +} + +func (s *postSaveHookStore) ListSummaries(ctx context.Context) ([]agentsession.Summary, error) { + return s.base.ListSummaries(ctx) +} + +func TestResolveCompactProviderSelectionResolveErrorBranch(t *testing.T) { + t.Parallel() + + manager := newRuntimeConfigManager(t) + cfg := manager.Get() + providerCfg, err := cfg.ProviderByName(cfg.SelectedProvider) + if err != nil { + t.Fatalf("ProviderByName() error = %v", err) + } + apiEnv := providerCfg.APIKeyEnv + restoreRuntimeEnv(t, apiEnv) + _ = os.Unsetenv(apiEnv) + + session := agentsession.Session{Provider: cfg.SelectedProvider, Model: "m1"} + if _, _, err := resolveCompactProviderSelection(session, cfg); err == nil { + t.Fatalf("expected resolve API key error") + } +} + +func TestGenerateStreamingMessageHooksAndContextDoneBranches(t *testing.T) { + t.Parallel() + + onDoneCalled := false + providerOK := &scriptedProvider{streams: [][]providertypes.StreamEvent{{providertypes.NewTextDeltaStreamEvent("ok")}}} + outcome := generateStreamingMessage(context.Background(), providerOK, providertypes.GenerateRequest{}, streaming.Hooks{ + OnMessageDone: func(payload providertypes.MessageDonePayload) { + onDoneCalled = true + }, + }) + if outcome.err != nil { + t.Fatalf("generateStreamingMessage() error = %v", outcome.err) + } + if !onDoneCalled { + t.Fatalf("expected OnMessageDone hook to be invoked") + } + + ctx, cancel := context.WithCancel(context.Background()) + providerWait := &scriptedProvider{chatFn: func(ctx context.Context, req providertypes.GenerateRequest, events chan<- providertypes.StreamEvent) error { + <-ctx.Done() + return ctx.Err() + }} + go func() { + time.Sleep(20 * time.Millisecond) + cancel() + }() + outcome = generateStreamingMessage(ctx, providerWait, providertypes.GenerateRequest{}, streaming.Hooks{}) + if !errors.Is(outcome.err, context.Canceled) { + t.Fatalf("expected context canceled, got %v", outcome.err) + } +} + +func TestPrepareTurnSnapshotErrorBranches(t *testing.T) { + t.Parallel() + + manager := newRuntimeConfigManager(t) + service := &Service{ + configManager: manager, + contextBuilder: &stubContextBuilder{buildFn: func(ctx context.Context, input agentcontext.BuildInput) (agentcontext.BuildResult, error) { + return agentcontext.BuildResult{}, errors.New("build failed") + }}, + toolManager: &stubToolManager{}, + } + state := newRunState("run-snapshot", newRuntimeSession("session-snapshot")) + if _, _, err := service.prepareTurnSnapshot(context.Background(), &state); err == nil { + t.Fatalf("expected build error") + } + + if err := manager.Update(context.Background(), func(cfg *config.Config) error { + cfg.SelectedProvider = "" + return nil + }); err != nil { + t.Fatalf("update config: %v", err) + } + service.contextBuilder = &stubContextBuilder{buildFn: func(ctx context.Context, input agentcontext.BuildInput) (agentcontext.BuildResult, error) { + return agentcontext.BuildResult{Messages: input.Messages}, nil + }} + if _, _, err := service.prepareTurnSnapshot(context.Background(), &state); err == nil { + t.Fatalf("expected resolve selected provider error") + } +} + +func TestApplyCompactForStateStrictErrorBranch(t *testing.T) { + t.Parallel() + + service := &Service{events: make(chan RuntimeEvent, 8)} + state := newRunState("run-apply-compact", newRuntimeSession("session-apply-compact")) + if _, err := service.applyCompactForState(context.Background(), &state, config.Config{}, contextcompact.ModeManual, compactErrorStrict); err == nil { + t.Fatalf("expected strict compact error") + } +} + +func TestExecuteToolCallWithPermissionRemainingBranches(t *testing.T) { + t.Parallel() + + askErr := permissionDecisionAskError(t) + + t.Run("emit chunk reads canceled context", func(t *testing.T) { + t.Parallel() + manager := &callbackToolManager{executeFn: func(ctx context.Context, input tools.ToolCallInput) (tools.ToolResult, error) { + if err := input.EmitChunk([]byte("x")); !errors.Is(err, context.Canceled) { + t.Fatalf("expected canceled emit chunk, got %v", err) + } + return tools.ToolResult{Name: input.Name}, context.Canceled + }} + service := &Service{toolManager: manager, approvalBroker: approvalflow.NewBroker(), events: make(chan RuntimeEvent, 8)} + ctx, cancel := context.WithCancel(context.Background()) + cancel() + _, _ = service.executeToolCallWithPermission(ctx, permissionExecutionInput{RunID: "r", SessionID: "s", Call: providertypes.ToolCall{ID: "c", Name: "filesystem_read_file"}, ToolTimeout: time.Second}) + }) + + t.Run("await permission decision error path", func(t *testing.T) { + t.Parallel() + manager := &callbackToolManager{executeFn: func(ctx context.Context, input tools.ToolCallInput) (tools.ToolResult, error) { + return tools.ToolResult{Name: input.Name}, askErr + }} + service := &Service{toolManager: manager, approvalBroker: nil, events: make(chan RuntimeEvent, 8)} + _, err := service.executeToolCallWithPermission(context.Background(), permissionExecutionInput{RunID: "r", SessionID: "s", Call: providertypes.ToolCall{ID: "c", Name: "filesystem_read_file"}, ToolTimeout: time.Second}) + if err == nil { + t.Fatalf("expected await permission error") + } + }) + + t.Run("invalid decision remember scope error", func(t *testing.T) { + t.Parallel() + manager := &callbackToolManager{executeFn: func(ctx context.Context, input tools.ToolCallInput) (tools.ToolResult, error) { + return tools.ToolResult{Name: input.Name}, askErr + }} + service := &Service{toolManager: manager, approvalBroker: approvalflow.NewBroker(), events: make(chan RuntimeEvent, 16)} + go func() { + for evt := range service.events { + if evt.Type != EventPermissionRequest { + continue + } + payload := evt.Payload.(PermissionRequestPayload) + _ = service.approvalBroker.Resolve(payload.RequestID, approvalflow.Decision("invalid")) + return + } + }() + _, err := service.executeToolCallWithPermission(context.Background(), permissionExecutionInput{RunID: "r", SessionID: "s", Call: providertypes.ToolCall{ID: "c", Name: "filesystem_read_file"}, ToolTimeout: time.Second}) + if err == nil { + t.Fatalf("expected invalid decision error") + } + }) + + t.Run("reject remember decision error", func(t *testing.T) { + t.Parallel() + manager := &callbackToolManager{ + executeFn: func(ctx context.Context, input tools.ToolCallInput) (tools.ToolResult, error) { + return tools.ToolResult{Name: input.Name}, askErr + }, + rememberErr: errors.New("remember failed"), + } + service := &Service{toolManager: manager, approvalBroker: approvalflow.NewBroker(), events: make(chan RuntimeEvent, 16)} + go resolveDecisionFromEvent(service, approvalflow.DecisionReject) + _, err := service.executeToolCallWithPermission(context.Background(), permissionExecutionInput{RunID: "r", SessionID: "s", Call: providertypes.ToolCall{ID: "c", Name: "filesystem_read_file"}, ToolTimeout: time.Second}) + if err == nil { + t.Fatalf("expected remember reject error") + } + }) + + t.Run("allow remember decision error", func(t *testing.T) { + t.Parallel() + manager := &callbackToolManager{ + executeFn: func(ctx context.Context, input tools.ToolCallInput) (tools.ToolResult, error) { + return tools.ToolResult{Name: input.Name}, askErr + }, + rememberErr: errors.New("remember failed"), + } + service := &Service{toolManager: manager, approvalBroker: approvalflow.NewBroker(), events: make(chan RuntimeEvent, 16)} + go resolveDecisionFromEvent(service, approvalflow.DecisionAllowSession) + _, err := service.executeToolCallWithPermission(context.Background(), permissionExecutionInput{RunID: "r", SessionID: "s", Call: providertypes.ToolCall{ID: "c", Name: "filesystem_read_file"}, ToolTimeout: time.Second}) + if err == nil { + t.Fatalf("expected remember allow error") + } + }) + + t.Run("non-ask deny emits fallback reason", func(t *testing.T) { + t.Parallel() + denyErr := permissionDecisionDenyError(t) + manager := &callbackToolManager{executeFn: func(ctx context.Context, input tools.ToolCallInput) (tools.ToolResult, error) { + return tools.ToolResult{Name: input.Name}, denyErr + }} + service := &Service{toolManager: manager, approvalBroker: approvalflow.NewBroker(), events: make(chan RuntimeEvent, 16)} + _, err := service.executeToolCallWithPermission(context.Background(), permissionExecutionInput{RunID: "r", SessionID: "s", Call: providertypes.ToolCall{ID: "c", Name: "filesystem_read_file"}, ToolTimeout: time.Second}) + if err == nil { + t.Fatalf("expected deny error") + } + events := collectRuntimeEvents(service.events) + found := false + for _, evt := range events { + if evt.Type != EventPermissionResolved { + continue + } + payload := evt.Payload.(PermissionResolvedPayload) + if payload.Reason != "" { + found = true + } + } + if !found { + t.Fatalf("expected permission resolved event with fallback reason") + } + }) +} + +func resolveDecisionFromEvent(service *Service, decision approvalflow.Decision) { + for evt := range service.events { + if evt.Type != EventPermissionRequest { + continue + } + payload := evt.Payload.(PermissionRequestPayload) + _ = service.approvalBroker.Resolve(payload.RequestID, decision) + return + } +} + +func permissionDecisionDenyError(t *testing.T) *tools.PermissionDecisionError { + t.Helper() + + registry := tools.NewRegistry() + registry.Register(&stubTool{name: "filesystem_read_file", content: "ok"}) + engine, err := security.NewStaticGateway(security.DecisionDeny, nil) + if err != nil { + t.Fatalf("new static gateway: %v", err) + } + manager, err := tools.NewManager(registry, engine, nil) + if err != nil { + t.Fatalf("new tool manager: %v", err) + } + _, execErr := manager.Execute(context.Background(), tools.ToolCallInput{ + ID: "call-deny", + Name: "filesystem_read_file", + Arguments: []byte(`{"path":"README.md"}`), + Workdir: t.TempDir(), + SessionID: "session-deny", + }) + var permissionErr *tools.PermissionDecisionError + if !errors.As(execErr, &permissionErr) { + t.Fatalf("expected PermissionDecisionError, got %v", execErr) + } + return permissionErr +} + +func TestExecuteAssistantToolCallsRemainingBranches(t *testing.T) { + t.Parallel() + + t.Run("returns ctx error before tool start", func(t *testing.T) { + t.Parallel() + service := &Service{events: make(chan RuntimeEvent, 8), approvalBroker: approvalflow.NewBroker()} + ctx, cancel := context.WithCancel(context.Background()) + cancel() + state := newRunState("run", newRuntimeSession("session-top-cancel")) + err := service.executeAssistantToolCalls(ctx, &state, turnSnapshot{}, providertypes.Message{ToolCalls: []providertypes.ToolCall{{ID: "c", Name: "filesystem_read_file"}}}) + if !errors.Is(err, context.Canceled) { + t.Fatalf("expected context.Canceled, got %v", err) + } + }) + + t.Run("successful tool but ctx canceled right after execution", func(t *testing.T) { + t.Parallel() + ctx, cancel := context.WithCancel(context.Background()) + manager := &callbackToolManager{executeFn: func(ctx context.Context, input tools.ToolCallInput) (tools.ToolResult, error) { + cancel() + return tools.ToolResult{Name: input.Name, Content: "ok"}, nil + }} + store := newMemoryStore() + session := newRuntimeSession("session-mid-cancel") + store.sessions[session.ID] = cloneSession(session) + service := &Service{events: make(chan RuntimeEvent, 8), approvalBroker: approvalflow.NewBroker(), toolManager: manager, sessionStore: store} + state := newRunState("run", session) + err := service.executeAssistantToolCalls(ctx, &state, turnSnapshot{toolTimeout: time.Second}, providertypes.Message{ToolCalls: []providertypes.ToolCall{{ID: "c", Name: "filesystem_read_file"}}}) + if !errors.Is(err, context.Canceled) { + t.Fatalf("expected context.Canceled, got %v", err) + } + }) + + t.Run("ctx canceled after save when execErr nil", func(t *testing.T) { + t.Parallel() + ctx, cancel := context.WithCancel(context.Background()) + manager := &callbackToolManager{executeFn: func(ctx context.Context, input tools.ToolCallInput) (tools.ToolResult, error) { + return tools.ToolResult{Name: input.Name, Content: "ok"}, nil + }} + base := newMemoryStore() + session := newRuntimeSession("session-after-save-cancel") + base.sessions[session.ID] = cloneSession(session) + store := &postSaveHookStore{base: base, saveHook: cancel} + service := &Service{events: make(chan RuntimeEvent, 8), approvalBroker: approvalflow.NewBroker(), toolManager: manager, sessionStore: store} + + state := newRunState("run", session) + err := service.executeAssistantToolCalls(ctx, &state, turnSnapshot{toolTimeout: time.Second}, providertypes.Message{ToolCalls: []providertypes.ToolCall{{ID: "c", Name: "filesystem_read_file"}}}) + if !errors.Is(err, context.Canceled) { + t.Fatalf("expected context.Canceled, got %v", err) + } + }) + + t.Run("ctx canceled after emit when execErr not nil", func(t *testing.T) { + t.Parallel() + ctx, cancel := context.WithCancel(context.Background()) + manager := &callbackToolManager{executeFn: func(ctx context.Context, input tools.ToolCallInput) (tools.ToolResult, error) { + return tools.ToolResult{Name: input.Name}, errors.New("tool failed") + }} + base := newMemoryStore() + session := newRuntimeSession("session-after-emit-cancel") + base.sessions[session.ID] = cloneSession(session) + store := &postSaveHookStore{base: base, saveHook: cancel} + service := &Service{events: make(chan RuntimeEvent, 8), approvalBroker: approvalflow.NewBroker(), toolManager: manager, sessionStore: store} + + state := newRunState("run", session) + err := service.executeAssistantToolCalls(ctx, &state, turnSnapshot{toolTimeout: time.Second}, providertypes.Message{ToolCalls: []providertypes.ToolCall{{ID: "c", Name: "filesystem_read_file"}}}) + if !errors.Is(err, context.Canceled) { + t.Fatalf("expected context.Canceled, got %v", err) + } + }) +} + +func TestRunAndProviderRetryRemainingBranches(t *testing.T) { + t.Parallel() + + t.Run("run exits when context canceled before snapshot", func(t *testing.T) { + t.Parallel() + ctx, cancel := context.WithCancel(context.Background()) + manager := newRuntimeConfigManager(t) + base := newMemoryStore() + existing := newRuntimeSession("session-cancel-before-snapshot") + existing.Workdir = manager.Get().Workdir + base.sessions[existing.ID] = cloneSession(existing) + store := &postSaveHookStore{base: base, saveHook: cancel} + service := NewWithFactory(manager, &stubToolManager{}, store, &scriptedProviderFactory{provider: &scriptedProvider{}}, &stubContextBuilder{}) + + err := service.Run(ctx, UserInput{RunID: "run-cancel-before-snapshot", SessionID: existing.ID, Content: "hello"}) + if !errors.Is(err, context.Canceled) { + t.Fatalf("expected context.Canceled, got %v", err) + } + }) + + t.Run("run returns assistant save error", func(t *testing.T) { + t.Parallel() + manager := newRuntimeConfigManager(t) + base := newMemoryStore() + existing := newRuntimeSession("session-assistant-save-fail") + existing.Workdir = manager.Get().Workdir + base.sessions[existing.ID] = cloneSession(existing) + store := &failingStore{Store: base, saveErr: errors.New("assistant save failed"), failOnSave: 2} + + service := NewWithFactory(manager, &stubToolManager{}, store, &scriptedProviderFactory{provider: &scriptedProvider{streams: [][]providertypes.StreamEvent{{providertypes.NewTextDeltaStreamEvent("answer")}}}}, &stubContextBuilder{}) + + err := service.Run(context.Background(), UserInput{RunID: "run-assistant-save-fail", SessionID: existing.ID, Content: "hello"}) + if err == nil || !containsError(err, "assistant save failed") { + t.Fatalf("expected assistant save error, got %v", err) + } + }) + + t.Run("callProviderWithRetry exits on context done during backoff wait", func(t *testing.T) { + t.Parallel() + ctx, cancel := context.WithCancel(context.Background()) + firstCallDone := make(chan struct{}, 1) + providerRetry := &scriptedProvider{chatFn: func(ctx context.Context, req providertypes.GenerateRequest, events chan<- providertypes.StreamEvent) error { + select { + case firstCallDone <- struct{}{}: + default: + } + return &provider.ProviderError{StatusCode: 500, Code: provider.ErrorCodeServer, Message: "retry", Retryable: true} + }} + go func() { + <-firstCallDone + cancel() + }() + service := &Service{providerFactory: &scriptedProviderFactory{provider: providerRetry}, events: make(chan RuntimeEvent, 8)} + state := newRunState("run-retry-backoff", newRuntimeSession("session-retry-backoff")) + _, err := service.callProviderWithRetry(ctx, &state, turnSnapshot{providerConfig: provider.RuntimeConfig{Name: "x"}}) + if !errors.Is(err, context.Canceled) { + t.Fatalf("expected context.Canceled, got %v", err) + } + }) + + t.Run("callProviderWithRetry checks ctx after retryable stream error", func(t *testing.T) { + t.Parallel() + ctx, cancel := context.WithCancel(context.Background()) + providerRetry := &scriptedProvider{chatFn: func(ctx context.Context, req providertypes.GenerateRequest, events chan<- providertypes.StreamEvent) error { + cancel() + return &provider.ProviderError{StatusCode: 500, Code: provider.ErrorCodeServer, Message: "retry", Retryable: true} + }} + service := &Service{providerFactory: &scriptedProviderFactory{provider: providerRetry}, events: make(chan RuntimeEvent, 8)} + state := newRunState("run-retry-ctx-check", newRuntimeSession("session-retry-ctx-check")) + _, err := service.callProviderWithRetry(ctx, &state, turnSnapshot{providerConfig: provider.RuntimeConfig{Name: "x"}}) + if !errors.Is(err, context.Canceled) { + t.Fatalf("expected context.Canceled, got %v", err) + } + }) +} diff --git a/internal/runtime/runtime_test.go b/internal/runtime/runtime_test.go index e92f7499..b0667462 100644 --- a/internal/runtime/runtime_test.go +++ b/internal/runtime/runtime_test.go @@ -14,7 +14,9 @@ import ( agentcontext "neo-code/internal/context" contextcompact "neo-code/internal/context/compact" "neo-code/internal/provider" + "neo-code/internal/provider/streaming" providertypes "neo-code/internal/provider/types" + approvalflow "neo-code/internal/runtime/approval" "neo-code/internal/security" agentsession "neo-code/internal/session" "neo-code/internal/tools" @@ -1134,7 +1136,7 @@ waitRequest: if err := service.ResolvePermission(context.Background(), PermissionResolutionInput{ RequestID: requestPayload.RequestID, - Decision: PermissionResolutionAllowSession, + Decision: approvalflow.DecisionAllowSession, }); err != nil { t.Fatalf("ResolvePermission() error = %v", err) } @@ -2051,9 +2053,10 @@ func TestServiceCompactManualFailureReturnsError(t *testing.T) { } func TestServiceCompactUsesSessionProviderAndModelWhenPresent(t *testing.T) { - t.Parallel() - geminiEnv := runtimeTestAPIKeyEnv(t) + "_GEMINI" + tempHome := t.TempDir() + setRuntimeEnv(t, "USERPROFILE", tempHome) + setRuntimeEnv(t, "HOME", tempHome) manager := newRuntimeConfigManagerWithProviderEnvs(t, map[string]string{ config.GeminiName: geminiEnv, }) @@ -2125,9 +2128,10 @@ func TestServiceCompactUsesSessionProviderAndModelWhenPresent(t *testing.T) { } func TestServiceCompactFallsBackToCurrentProviderWhenSessionMetadataMissing(t *testing.T) { - t.Parallel() - geminiEnv := runtimeTestAPIKeyEnv(t) + "_GEMINI" + tempHome := t.TempDir() + setRuntimeEnv(t, "USERPROFILE", tempHome) + setRuntimeEnv(t, "HOME", tempHome) manager := newRuntimeConfigManagerWithProviderEnvs(t, map[string]string{ config.GeminiName: geminiEnv, }) @@ -2579,6 +2583,14 @@ func setRuntimeProviderEnv(t *testing.T, key string, value string) { } } +func setRuntimeEnv(t *testing.T, key string, value string) { + t.Helper() + restoreRuntimeEnv(t, key) + if err := os.Setenv(key, value); err != nil { + t.Fatalf("set env %s: %v", key, err) + } +} + func restoreRuntimeEnv(t *testing.T, key string) { t.Helper() value, ok := os.LookupEnv(key) @@ -2696,10 +2708,10 @@ func containsError(err error, target string) bool { func TestWorkdirHelperFunctions(t *testing.T) { t.Run("effectiveSessionWorkdir prefers session value", func(t *testing.T) { - if got := effectiveSessionWorkdir(" /session ", "/default"); got != "/session" { + if got := agentsession.EffectiveWorkdir(" /session ", "/default"); got != "/session" { t.Fatalf("expected session workdir, got %q", got) } - if got := effectiveSessionWorkdir("", " /default "); got != "/default" { + if got := agentsession.EffectiveWorkdir("", " /default "); got != "/default" { t.Fatalf("expected default workdir, got %q", got) } }) @@ -2737,7 +2749,7 @@ func TestWorkdirHelperFunctions(t *testing.T) { if err := os.WriteFile(filePath, []byte("x"), 0o644); err != nil { t.Fatalf("write file: %v", err) } - _, err = normalizeExistingWorkdir(filePath) + _, err = agentsession.ResolveExistingDir(filePath) if err == nil || !containsError(err, "is not a directory") { t.Fatalf("expected non-directory error, got %v", err) } @@ -2813,58 +2825,14 @@ func TestProviderRetryBackoff(t *testing.T) { } } -func TestPermissionEventViewPayloadMapping(t *testing.T) { - t.Parallel() - - view := permissionEventView{ - toolName: "webfetch", - actionType: string(security.ActionTypeRead), - operation: "fetch", - targetType: string(security.TargetTypeURL), - target: "https://example.com", - decision: "ask", - reason: "need approval", - ruleID: "rule-1", - scope: string(tools.SessionPermissionScopeAlways), - resolvedAs: "rejected", - } - - requestPayload := view.toRequestPayload() - if requestPayload.ToolName != view.toolName || - requestPayload.ActionType != view.actionType || - requestPayload.Operation != view.operation || - requestPayload.TargetType != view.targetType || - requestPayload.Target != view.target || - requestPayload.Decision != view.decision || - requestPayload.Reason != view.reason || - requestPayload.RuleID != view.ruleID || - requestPayload.RememberScope != view.scope { - t.Fatalf("unexpected request payload: %+v", requestPayload) - } - - resolvedPayload := view.toResolvedPayload() - if resolvedPayload.ToolName != view.toolName || - resolvedPayload.ActionType != view.actionType || - resolvedPayload.Operation != view.operation || - resolvedPayload.TargetType != view.targetType || - resolvedPayload.Target != view.target || - resolvedPayload.Decision != view.decision || - resolvedPayload.Reason != view.reason || - resolvedPayload.RuleID != view.ruleID || - resolvedPayload.RememberScope != view.scope || - resolvedPayload.ResolvedAs != view.resolvedAs { - t.Fatalf("unexpected resolved payload: %+v", resolvedPayload) - } -} - func TestStreamAccumulatorBuildMessageRejectsMissingToolName(t *testing.T) { t.Parallel() - acc := newStreamAccumulator() - acc.accumulateToolCallStart(0, "call-1", "") - acc.accumulateToolCallDelta(0, "call-1", "{}") + acc := streaming.NewAccumulator() + acc.AccumulateToolCallStart(0, "call-1", "") + acc.AccumulateToolCallDelta(0, "call-1", "{}") - _, err := acc.buildMessage() + _, err := acc.BuildMessage() if err == nil || !containsError(err, "without name") { t.Fatalf("expected missing tool name error, got %v", err) } @@ -2937,36 +2905,30 @@ func TestServiceRunFailsWhenAssistantSaveFails(t *testing.T) { func TestHandleProviderStreamEventErrorBranches(t *testing.T) { t.Parallel() - acc := newStreamAccumulator() + acc := streaming.NewAccumulator() - err := handleProviderStreamEvent( + err := streaming.HandleEvent( providertypes.StreamEvent{Type: providertypes.StreamEventToolCallStart}, acc, - nil, - nil, - nil, + streaming.Hooks{}, ) if err == nil || !containsError(err, "tool_call_start event payload is nil") { t.Fatalf("expected tool_call_start payload error, got %v", err) } - err = handleProviderStreamEvent( + err = streaming.HandleEvent( providertypes.StreamEvent{Type: providertypes.StreamEventToolCallDelta}, acc, - nil, - nil, - nil, + streaming.Hooks{}, ) if err == nil || !containsError(err, "tool_call_delta event payload is nil") { t.Fatalf("expected tool_call_delta payload error, got %v", err) } - err = handleProviderStreamEvent( + err = streaming.HandleEvent( providertypes.StreamEvent{Type: providertypes.StreamEventMessageDone}, acc, - nil, - nil, - nil, + streaming.Hooks{}, ) if err == nil || !containsError(err, "message_done event payload is nil") { t.Fatalf("expected message_done payload error, got %v", err) @@ -3013,15 +2975,20 @@ func TestCallProviderWithRetryReturnsCombinedForwardError(t *testing.T) { } service := NewWithFactory(manager, registry, store, &scriptedProviderFactory{provider: scripted}, nil) - _, err := service.callProviderWithRetry( - context.Background(), - "run-forward-error", - "session-forward-error", - providertypes.GenerateRequest{ + state := newRunState("run-forward-error", agentsession.Session{ID: "session-forward-error"}) + snapshot := turnSnapshot{ + providerConfig: provider.RuntimeConfig{}, + request: providertypes.GenerateRequest{ Model: "test-model", SystemPrompt: "prompt", Messages: []providertypes.Message{{Role: providertypes.RoleUser, Content: "hello"}}, }, + } + + _, err := service.callProviderWithRetry( + context.Background(), + &state, + snapshot, ) if err == nil || !containsError(err, "provider stream handling failed after provider error") { t.Fatalf("expected combined forward/provider error, got %v", err) @@ -3194,9 +3161,9 @@ func TestServiceRunAutoCompactsAndResetsSessionTokens(t *testing.T) { builder := &stubContextBuilder{ buildFn: func(ctx context.Context, input agentcontext.BuildInput) (agentcontext.BuildResult, error) { return agentcontext.BuildResult{ - SystemPrompt: "auto compact prompt", - Messages: append([]providertypes.Message(nil), input.Messages...), - ShouldAutoCompact: input.Metadata.SessionInputTokens >= input.Compact.AutoCompactThreshold, + SystemPrompt: "auto compact prompt", + Messages: append([]providertypes.Message(nil), input.Messages...), + AutoCompactSuggested: input.Metadata.SessionInputTokens >= input.Compact.AutoCompactThreshold, }, nil }, } @@ -3282,13 +3249,6 @@ func TestServiceRunAutoCompactsAndResetsSessionTokens(t *testing.T) { t.Fatalf("expected first provider request to use compacted latest answer, got %+v", scripted.requests[0].Messages) } - if service.sessionInputTokens != 0 { - t.Fatalf("expected service input tokens to reset, got %d", service.sessionInputTokens) - } - if service.sessionOutputTokens != 0 { - t.Fatalf("expected service output tokens to reset, got %d", service.sessionOutputTokens) - } - saved, err := store.Load(context.Background(), session.ID) if err != nil { t.Fatalf("load compacted session: %v", err) @@ -3319,9 +3279,9 @@ func TestServiceRunAutoCompactsAndResetsSessionTokens(t *testing.T) { if event.Type != EventCompactDone { continue } - payload, ok := event.Payload.(CompactDonePayload) + payload, ok := event.Payload.(CompactResult) if !ok { - t.Fatalf("expected CompactDonePayload, got %T", event.Payload) + t.Fatalf("expected CompactResult, got %T", event.Payload) } if payload.TriggerMode != string(contextcompact.ModeAuto) { t.Fatalf("expected trigger mode %q, got %q", contextcompact.ModeAuto, payload.TriggerMode) @@ -3361,9 +3321,9 @@ func TestServiceRunAutoCompactNoopDoesNotDisableReactiveRetry(t *testing.T) { builder := &stubContextBuilder{ buildFn: func(ctx context.Context, input agentcontext.BuildInput) (agentcontext.BuildResult, error) { return agentcontext.BuildResult{ - SystemPrompt: "auto compact prompt", - Messages: append([]providertypes.Message(nil), input.Messages...), - ShouldAutoCompact: input.Metadata.SessionInputTokens >= input.Compact.AutoCompactThreshold, + SystemPrompt: "auto compact prompt", + Messages: append([]providertypes.Message(nil), input.Messages...), + AutoCompactSuggested: input.Metadata.SessionInputTokens >= input.Compact.AutoCompactThreshold, }, nil }, } @@ -3579,9 +3539,9 @@ func TestServiceRunReactivelyCompactsOnContextTooLong(t *testing.T) { if event.Type != EventCompactDone { continue } - payload, ok := event.Payload.(CompactDonePayload) + payload, ok := event.Payload.(CompactResult) if !ok { - t.Fatalf("expected CompactDonePayload, got %T", event.Payload) + t.Fatalf("expected CompactResult, got %T", event.Payload) } if payload.TriggerMode != string(contextcompact.ModeReactive) { t.Fatalf("expected trigger mode %q, got %q", contextcompact.ModeReactive, payload.TriggerMode) @@ -3794,51 +3754,42 @@ func TestServiceRunDoesNotReactiveCompactOnPlainTextTokenThrottle(t *testing.T) func TestRestoreSessionTokens(t *testing.T) { t.Parallel() - service := &Service{ - events: make(chan RuntimeEvent, 128), - } session := agentsession.Session{ TokenInputTotal: 500, TokenOutputTotal: 200, } - service.restoreSessionTokens(session) + state := newRunState("", session) - if service.sessionInputTokens != 500 { - t.Fatalf("expected sessionInputTokens == 500, got %d", service.sessionInputTokens) + if state.tokenInputTotal != 500 { + t.Fatalf("expected sessionInputTokens == 500, got %d", state.tokenInputTotal) } - if service.sessionOutputTokens != 200 { - t.Fatalf("expected sessionOutputTokens == 200, got %d", service.sessionOutputTokens) + if state.tokenOutputTotal != 200 { + t.Fatalf("expected sessionOutputTokens == 200, got %d", state.tokenOutputTotal) } } func TestRestoreSessionTokensNewSession(t *testing.T) { t.Parallel() - service := &Service{ - events: make(chan RuntimeEvent, 128), - } session := agentsession.Session{ TokenInputTotal: 0, TokenOutputTotal: 0, } - service.restoreSessionTokens(session) + state := newRunState("", session) - if service.sessionInputTokens != 0 { - t.Fatalf("expected sessionInputTokens == 0, got %d", service.sessionInputTokens) + if state.tokenInputTotal != 0 { + t.Fatalf("expected sessionInputTokens == 0, got %d", state.tokenInputTotal) } - if service.sessionOutputTokens != 0 { - t.Fatalf("expected sessionOutputTokens == 0, got %d", service.sessionOutputTokens) + if state.tokenOutputTotal != 0 { + t.Fatalf("expected sessionOutputTokens == 0, got %d", state.tokenOutputTotal) } } func TestAutoCompactThresholdEnabled(t *testing.T) { t.Parallel() - service := &Service{ - events: make(chan RuntimeEvent, 128), - } cfg := config.Config{ Context: config.ContextConfig{ AutoCompact: config.AutoCompactConfig{ @@ -3848,7 +3799,7 @@ func TestAutoCompactThresholdEnabled(t *testing.T) { }, } - threshold := service.autoCompactThreshold(cfg) + threshold := autoCompactThreshold(cfg) if threshold != 50000 { t.Fatalf("expected threshold == 50000, got %d", threshold) } @@ -3857,9 +3808,6 @@ func TestAutoCompactThresholdEnabled(t *testing.T) { func TestAutoCompactThresholdDisabled(t *testing.T) { t.Parallel() - service := &Service{ - events: make(chan RuntimeEvent, 128), - } cfg := config.Config{ Context: config.ContextConfig{ AutoCompact: config.AutoCompactConfig{ @@ -3869,7 +3817,7 @@ func TestAutoCompactThresholdDisabled(t *testing.T) { }, } - threshold := service.autoCompactThreshold(cfg) + threshold := autoCompactThreshold(cfg) if threshold != 0 { t.Fatalf("expected threshold == 0, got %d", threshold) } @@ -3878,9 +3826,6 @@ func TestAutoCompactThresholdDisabled(t *testing.T) { func TestAutoCompactThresholdZeroValue(t *testing.T) { t.Parallel() - service := &Service{ - events: make(chan RuntimeEvent, 128), - } cfg := config.Config{ Context: config.ContextConfig{ AutoCompact: config.AutoCompactConfig{ @@ -3890,7 +3835,7 @@ func TestAutoCompactThresholdZeroValue(t *testing.T) { }, } - threshold := service.autoCompactThreshold(cfg) + threshold := autoCompactThreshold(cfg) if threshold != 0 { t.Fatalf("expected threshold == 0, got %d", threshold) } @@ -3900,10 +3845,9 @@ func TestTokenUsageRecordedOnMessageDone(t *testing.T) { t.Parallel() service := &Service{ - events: make(chan RuntimeEvent, 128), - sessionInputTokens: 0, - sessionOutputTokens: 0, + events: make(chan RuntimeEvent, 128), } + state := runState{} events := collectRuntimeEvents(service.Events()) @@ -3913,35 +3857,32 @@ func TestTokenUsageRecordedOnMessageDone(t *testing.T) { OutputTokens: 50, }) - // Handle the event with an onMessageDone callback that mimics forwardProviderEvents - err := handleProviderStreamEvent( + // 使用与运行时相同的流式事件处理器验证 usage 累积行为。 + err := streaming.HandleEvent( messageDoneEvent, nil, - nil, - nil, - func(payload providertypes.MessageDonePayload) { + streaming.Hooks{OnMessageDone: func(payload providertypes.MessageDonePayload) { if payload.Usage != nil { - service.sessionInputTokens += payload.Usage.InputTokens - service.sessionOutputTokens += payload.Usage.OutputTokens + state.recordUsage(payload.Usage.InputTokens, payload.Usage.OutputTokens) service.emit(context.Background(), EventTokenUsage, "test-run-id", "test-session-id", TokenUsagePayload{ InputTokens: payload.Usage.InputTokens, OutputTokens: payload.Usage.OutputTokens, - SessionInputTokens: service.sessionInputTokens, - SessionOutputTokens: service.sessionOutputTokens, + SessionInputTokens: state.tokenInputTotal, + SessionOutputTokens: state.tokenOutputTotal, }) } - }, + }}, ) if err != nil { - t.Fatalf("handleProviderStreamEvent error = %v", err) + t.Fatalf("streaming.HandleEvent error = %v", err) } // Verify the service counters are updated - if service.sessionInputTokens != 100 { - t.Fatalf("expected sessionInputTokens == 100, got %d", service.sessionInputTokens) + if state.tokenInputTotal != 100 { + t.Fatalf("expected sessionInputTokens == 100, got %d", state.tokenInputTotal) } - if service.sessionOutputTokens != 50 { - t.Fatalf("expected sessionOutputTokens == 50, got %d", service.sessionOutputTokens) + if state.tokenOutputTotal != 50 { + t.Fatalf("expected sessionOutputTokens == 50, got %d", state.tokenOutputTotal) } // Verify EventTokenUsage was emitted with correct payload diff --git a/internal/runtime/session_mutation.go b/internal/runtime/session_mutation.go new file mode 100644 index 00000000..9e3ab972 --- /dev/null +++ b/internal/runtime/session_mutation.go @@ -0,0 +1,67 @@ +package runtime + +import ( + "context" + "strings" + + providertypes "neo-code/internal/provider/types" + "neo-code/internal/tools" +) + +// appendUserMessageAndSave 将用户消息追加到会话并立即持久化。 +func (s *Service) appendUserMessageAndSave(ctx context.Context, state *runState, content string) error { + message := providertypes.Message{ + Role: providertypes.RoleUser, + Content: content, + } + state.session.Messages = append(state.session.Messages, message) + state.touchSession() + if err := s.sessionStore.Save(ctx, &state.session); err != nil { + return err + } + s.emit(ctx, EventUserMessage, state.runID, state.session.ID, message) + return nil +} + +// appendAssistantMessageAndSave 将 assistant 消息和本轮模型元数据写回会话。 +func (s *Service) appendAssistantMessageAndSave( + ctx context.Context, + state *runState, + snapshot turnSnapshot, + assistant providertypes.Message, +) error { + metadataChanged := state.session.Provider != snapshot.providerConfig.Name || state.session.Model != snapshot.model + state.session.Provider = snapshot.providerConfig.Name + state.session.Model = snapshot.model + state.syncSessionTokenTotals() + + if strings.TrimSpace(assistant.Content) != "" || len(assistant.ToolCalls) > 0 { + state.session.Messages = append(state.session.Messages, assistant) + state.touchSession() + return s.sessionStore.Save(ctx, &state.session) + } + if metadataChanged { + state.touchSession() + return s.sessionStore.Save(ctx, &state.session) + } + return nil +} + +// appendToolMessageAndSave 将工具原始结果写回会话，避免污染持久化对话内容。 +func (s *Service) appendToolMessageAndSave( + ctx context.Context, + state *runState, + call providertypes.ToolCall, + result tools.ToolResult, +) error { + toolMessage := providertypes.Message{ + Role: providertypes.RoleTool, + Content: result.Content, + ToolCallID: call.ID, + IsError: result.IsError, + ToolMetadata: tools.SanitizeToolMetadata(result.Name, result.Metadata), + } + state.session.Messages = append(state.session.Messages, toolMessage) + state.touchSession() + return s.sessionStore.Save(ctx, &state.session) +} diff --git a/internal/runtime/state.go b/internal/runtime/state.go new file mode 100644 index 00000000..7a795961 --- /dev/null +++ b/internal/runtime/state.go @@ -0,0 +1,84 @@ +package runtime + +import ( + "time" + + "neo-code/internal/config" + "neo-code/internal/provider" + providertypes "neo-code/internal/provider/types" + agentsession "neo-code/internal/session" +) + +// runState 汇总单次 Run 生命周期内会变化的会话与计量状态。 +type runState struct { + runID string + session agentsession.Session + tokenInputTotal int + tokenOutputTotal int + compactApplied bool + reactiveCompactUsed bool +} + +// newRunState 基于持久化会话创建一次运行的内存状态镜像。 +func newRunState(runID string, session agentsession.Session) runState { + return runState{ + runID: runID, + session: session, + tokenInputTotal: session.TokenInputTotal, + tokenOutputTotal: session.TokenOutputTotal, + } +} + +// syncSessionTokenTotals 将运行期 token 计数同步回会话对象。 +func (s *runState) syncSessionTokenTotals() { + if s == nil { + return + } + s.session.TokenInputTotal = s.tokenInputTotal + s.session.TokenOutputTotal = s.tokenOutputTotal +} + +// recordUsage 累加本轮 provider 返回的 token 使用量。 +func (s *runState) recordUsage(inputTokens int, outputTokens int) { + if s == nil { + return + } + s.tokenInputTotal += inputTokens + s.tokenOutputTotal += outputTokens +} + +// resetTokenTotals 在 compact 应用成功后清零当前运行的 token 账本。 +func (s *runState) resetTokenTotals() { + if s == nil { + return + } + s.tokenInputTotal = 0 + s.tokenOutputTotal = 0 + s.syncSessionTokenTotals() +} + +// touchSession 更新会话修改时间并同步最新 token 累计值。 +func (s *runState) touchSession() { + if s == nil { + return + } + s.syncSessionTokenTotals() + s.session.UpdatedAt = time.Now() +} + +// turnSnapshot 冻结单轮推理所需的配置、上下文与 provider 请求。 +type turnSnapshot struct { + config config.Config + providerConfig provider.RuntimeConfig + model string + workdir string + toolTimeout time.Duration + request providertypes.GenerateRequest +} + +// providerTurnResult 表示单轮 provider 调用成功后的结构化结果。 +type providerTurnResult struct { + assistant providertypes.Message + inputTokens int + outputTokens int +} diff --git a/internal/runtime/toolexec.go b/internal/runtime/toolexec.go new file mode 100644 index 00000000..37b389f5 --- /dev/null +++ b/internal/runtime/toolexec.go @@ -0,0 +1,61 @@ +package runtime + +import ( + "context" + "errors" + "strings" + + providertypes "neo-code/internal/provider/types" +) + +// executeAssistantToolCalls 顺序执行 assistant 返回的全部工具调用并回写结果。 +func (s *Service) executeAssistantToolCalls( + ctx context.Context, + state *runState, + snapshot turnSnapshot, + assistant providertypes.Message, +) error { + for _, call := range assistant.ToolCalls { + if err := ctx.Err(); err != nil { + return err + } + s.emit(ctx, EventToolStart, state.runID, state.session.ID, call) + + result, execErr := s.executeToolCallWithPermission(ctx, permissionExecutionInput{ + RunID: state.runID, + SessionID: state.session.ID, + Call: call, + Workdir: snapshot.workdir, + ToolTimeout: snapshot.toolTimeout, + }) + if errors.Is(execErr, context.Canceled) { + return execErr + } + if execErr == nil { + if err := ctx.Err(); err != nil { + return err + } + } + + if execErr != nil && strings.TrimSpace(result.Content) == "" { + result.Content = execErr.Error() + } + if err := s.appendToolMessageAndSave(ctx, state, call, result); err != nil { + if execErr != nil && errors.Is(err, context.Canceled) { + s.emit(ctx, EventToolResult, state.runID, state.session.ID, result) + } + return err + } + if err := ctx.Err(); err != nil && execErr == nil { + return err + } + + s.emit(ctx, EventToolResult, state.runID, state.session.ID, result) + if execErr != nil { + if err := ctx.Err(); err != nil { + return err + } + } + } + return nil +} diff --git a/internal/runtime/workdir_branch_test.go b/internal/runtime/workdir_branch_test.go index 089a3cd0..a1137140 100644 --- a/internal/runtime/workdir_branch_test.go +++ b/internal/runtime/workdir_branch_test.go @@ -41,7 +41,7 @@ func TestResolveWorkdirForSessionAndNormalizeErrors(t *testing.T) { t.Fatalf("expected empty workdir error, got %v", err) } - _, err = normalizeExistingWorkdir(filepath.Join(base, "missing")) + _, err = agentsession.ResolveExistingDir(filepath.Join(base, "missing")) if err == nil || !strings.Contains(err.Error(), "resolve workdir") { t.Fatalf("expected missing path error, got %v", err) } @@ -50,7 +50,7 @@ func TestResolveWorkdirForSessionAndNormalizeErrors(t *testing.T) { if err := os.WriteFile(filePath, []byte("x"), 0o644); err != nil { t.Fatalf("write file: %v", err) } - _, err = normalizeExistingWorkdir(filePath) + _, err = agentsession.ResolveExistingDir(filePath) if err == nil || !strings.Contains(err.Error(), "is not a directory") { t.Fatalf("expected non-directory error, got %v", err) } diff --git a/internal/tui/core/app/permission_prompt.go b/internal/tui/core/app/permission_prompt.go index 0bf31392..02fa1410 100644 --- a/internal/tui/core/app/permission_prompt.go +++ b/internal/tui/core/app/permission_prompt.go @@ -8,6 +8,7 @@ import ( "github.com/charmbracelet/lipgloss" agentruntime "neo-code/internal/runtime" + approvalflow "neo-code/internal/runtime/approval" ) // permissionPromptOption 表示权限审批面板中的一个可选项。 @@ -21,17 +22,17 @@ var permissionPromptOptions = []permissionPromptOption{ { Label: "Allow once", Hint: "Approve this request once", - Decision: agentruntime.PermissionResolutionAllowOnce, + Decision: approvalflow.DecisionAllowOnce, }, { Label: "Allow session", Hint: "Approve similar requests for this session", - Decision: agentruntime.PermissionResolutionAllowSession, + Decision: approvalflow.DecisionAllowSession, }, { Label: "Reject", Hint: "Reject this request", - Decision: agentruntime.PermissionResolutionReject, + Decision: approvalflow.DecisionReject, }, } @@ -66,11 +67,11 @@ func permissionPromptOptionAt(selected int) permissionPromptOption { func parsePermissionShortcut(input string) (agentruntime.PermissionResolutionDecision, bool) { switch strings.ToLower(strings.TrimSpace(input)) { case "y", "yes", "once": - return agentruntime.PermissionResolutionAllowOnce, true + return approvalflow.DecisionAllowOnce, true case "a", "always": - return agentruntime.PermissionResolutionAllowSession, true + return approvalflow.DecisionAllowSession, true case "n", "no", "reject", "deny": - return agentruntime.PermissionResolutionReject, true + return approvalflow.DecisionReject, true default: return "", false } diff --git a/internal/tui/core/app/permission_prompt_test.go b/internal/tui/core/app/permission_prompt_test.go index db85fa81..0faceb81 100644 --- a/internal/tui/core/app/permission_prompt_test.go +++ b/internal/tui/core/app/permission_prompt_test.go @@ -7,6 +7,7 @@ import ( "github.com/charmbracelet/bubbles/textarea" agentruntime "neo-code/internal/runtime" + approvalflow "neo-code/internal/runtime/approval" ) func TestNormalizePermissionPromptSelectionWrap(t *testing.T) { @@ -30,19 +31,19 @@ func TestNormalizePermissionPromptSelectionEmptyOptions(t *testing.T) { func TestPermissionPromptOptionAt(t *testing.T) { option := permissionPromptOptionAt(-1) - if option.Decision != agentruntime.PermissionResolutionReject { + if option.Decision != approvalflow.DecisionReject { t.Fatalf("expected wrapped option to be reject, got %q", option.Decision) } } func TestParsePermissionShortcut(t *testing.T) { tests := map[string]agentruntime.PermissionResolutionDecision{ - "y": agentruntime.PermissionResolutionAllowOnce, - "once": agentruntime.PermissionResolutionAllowOnce, - "a": agentruntime.PermissionResolutionAllowSession, - "always": agentruntime.PermissionResolutionAllowSession, - "n": agentruntime.PermissionResolutionReject, - "deny": agentruntime.PermissionResolutionReject, + "y": approvalflow.DecisionAllowOnce, + "once": approvalflow.DecisionAllowOnce, + "a": approvalflow.DecisionAllowSession, + "always": approvalflow.DecisionAllowSession, + "n": approvalflow.DecisionReject, + "deny": approvalflow.DecisionReject, } for input, want := range tests { got, ok := parsePermissionShortcut(input) diff --git a/internal/tui/core/app/update.go b/internal/tui/core/app/update.go index 1dea5be0..c117524f 100644 --- a/internal/tui/core/app/update.go +++ b/internal/tui/core/app/update.go @@ -18,6 +18,7 @@ import ( "neo-code/internal/memo" providertypes "neo-code/internal/provider/types" agentruntime "neo-code/internal/runtime" + approvalflow "neo-code/internal/runtime/approval" agentsession "neo-code/internal/session" "neo-code/internal/tools" tuistatus "neo-code/internal/tui/core/status" @@ -1023,7 +1024,7 @@ func runtimeEventPermissionRequestHandler(a *App, event agentruntime.RuntimeEven currentRequestID := strings.TrimSpace(a.pendingPermission.Request.RequestID) nextRequestID := strings.TrimSpace(payload.RequestID) if currentRequestID != "" && currentRequestID != nextRequestID && !a.pendingPermission.Submitting { - a.deferredEventCmd = runResolvePermission(a.runtime, currentRequestID, agentruntime.PermissionResolutionReject) + a.deferredEventCmd = runResolvePermission(a.runtime, currentRequestID, approvalflow.DecisionReject) a.appendActivity( "permission", "Auto-rejected superseded permission request", @@ -1083,7 +1084,7 @@ func (a *App) refreshPermissionPromptLayout() { // runtimeEventCompactDoneHandler 处理 compact 完成事件。 func runtimeEventCompactDoneHandler(a *App, event agentruntime.RuntimeEvent) bool { - payload, ok := event.Payload.(agentruntime.CompactDonePayload) + payload, ok := event.Payload.(agentruntime.CompactResult) if !ok { return false } diff --git a/internal/tui/core/app/update_permission_test.go b/internal/tui/core/app/update_permission_test.go index 6f42fc75..fa0ad3af 100644 --- a/internal/tui/core/app/update_permission_test.go +++ b/internal/tui/core/app/update_permission_test.go @@ -13,6 +13,7 @@ import ( tea "github.com/charmbracelet/bubbletea" agentruntime "neo-code/internal/runtime" + approvalflow "neo-code/internal/runtime/approval" agentsession "neo-code/internal/session" tuistate "neo-code/internal/tui/state" ) @@ -123,10 +124,10 @@ func TestUpdatePendingPermissionInputSelectAndSubmit(t *testing.T) { if !ok { t.Fatalf("expected permissionResolutionFinishedMsg, got %T", msg) } - if done.RequestID != "perm-1" || done.Decision != agentruntime.PermissionResolutionAllowOnce { + if done.RequestID != "perm-1" || done.Decision != approvalflow.DecisionAllowOnce { t.Fatalf("unexpected submitted decision: %+v", done) } - if runtime.lastResolved.Decision != agentruntime.PermissionResolutionAllowOnce { + if runtime.lastResolved.Decision != approvalflow.DecisionAllowOnce { t.Fatalf("runtime decision mismatch: %+v", runtime.lastResolved) } } @@ -156,7 +157,7 @@ func TestUpdatePendingPermissionInputShortcut(t *testing.T) { if !ok { t.Fatalf("expected permissionResolutionFinishedMsg, got %T", msg) } - if done.Decision != agentruntime.PermissionResolutionReject { + if done.Decision != approvalflow.DecisionReject { t.Fatalf("expected reject decision, got %q", done.Decision) } } @@ -176,7 +177,7 @@ func TestUpdatePendingPermissionInputSubmittingConsumesInput(t *testing.T) { func TestSubmitPermissionDecisionValidation(t *testing.T) { app := newPermissionTestApp(&permissionTestRuntime{}) - if cmd := app.submitPermissionDecision(agentruntime.PermissionResolutionAllowOnce); cmd != nil { + if cmd := app.submitPermissionDecision(approvalflow.DecisionAllowOnce); cmd != nil { t.Fatalf("expected nil cmd when no pending permission") } @@ -184,7 +185,7 @@ func TestSubmitPermissionDecisionValidation(t *testing.T) { Request: agentruntime.PermissionRequestPayload{RequestID: " "}, Selected: 0, } - if cmd := app.submitPermissionDecision(agentruntime.PermissionResolutionAllowOnce); cmd != nil { + if cmd := app.submitPermissionDecision(approvalflow.DecisionAllowOnce); cmd != nil { t.Fatalf("expected nil cmd for empty request id") } } @@ -237,7 +238,7 @@ func TestUpdatePermissionResolutionFinishedMessage(t *testing.T) { model, _ := app.Update(permissionResolutionFinishedMsg{ RequestID: "perm-5", - Decision: agentruntime.PermissionResolutionAllowOnce, + Decision: approvalflow.DecisionAllowOnce, Err: errors.New("network"), }) next := model.(App) @@ -259,7 +260,7 @@ func TestUpdatePermissionResolutionFinishedMessageSuccessClearsPendingPermission model, _ := app.Update(permissionResolutionFinishedMsg{ RequestID: "perm-5-success", - Decision: agentruntime.PermissionResolutionAllowOnce, + Decision: approvalflow.DecisionAllowOnce, }) next := model.(App) if next.pendingPermission != nil { @@ -313,10 +314,10 @@ func TestRuntimePermissionRequestHandlerAutoRejectsSupersededRequest(t *testing. if !ok { t.Fatalf("expected permissionResolutionFinishedMsg, got %T", msg) } - if done.RequestID != "perm-old" || done.Decision != agentruntime.PermissionResolutionReject { + if done.RequestID != "perm-old" || done.Decision != approvalflow.DecisionReject { t.Fatalf("unexpected auto-reject payload: %+v", done) } - if runtime.lastResolved.RequestID != "perm-old" || runtime.lastResolved.Decision != agentruntime.PermissionResolutionReject { + if runtime.lastResolved.RequestID != "perm-old" || runtime.lastResolved.Decision != approvalflow.DecisionReject { t.Fatalf("unexpected runtime resolve input: %+v", runtime.lastResolved) } } @@ -370,7 +371,7 @@ func TestHandleRuntimeEventQueuesDeferredCommand(t *testing.T) { if _, ok := batch[0]().(permissionResolutionFinishedMsg); !ok { t.Fatalf("expected deferred batch command to resolve permission") } - if runtime.lastResolved.RequestID != "perm-old" || runtime.lastResolved.Decision != agentruntime.PermissionResolutionReject { + if runtime.lastResolved.RequestID != "perm-old" || runtime.lastResolved.Decision != approvalflow.DecisionReject { t.Fatalf("expected deferred auto-reject to run, got %+v", runtime.lastResolved) } } @@ -394,7 +395,7 @@ func TestRuntimePermissionResolvedHandlerUsesExactRequestIDMatch(t *testing.T) { func TestRunResolvePermissionForwardsRuntimeError(t *testing.T) { runtime := &permissionTestRuntime{resolveErr: errors.New("resolve failed")} - cmd := runResolvePermission(runtime, "perm-7", agentruntime.PermissionResolutionReject) + cmd := runResolvePermission(runtime, "perm-7", approvalflow.DecisionReject) msg := cmd() done, ok := msg.(permissionResolutionFinishedMsg) if !ok { @@ -403,7 +404,7 @@ func TestRunResolvePermissionForwardsRuntimeError(t *testing.T) { if done.Err == nil || done.Err.Error() != "resolve failed" { t.Fatalf("expected forwarded resolve error, got %#v", done.Err) } - if runtime.lastResolved.RequestID != "perm-7" || runtime.lastResolved.Decision != agentruntime.PermissionResolutionReject { + if runtime.lastResolved.RequestID != "perm-7" || runtime.lastResolved.Decision != approvalflow.DecisionReject { t.Fatalf("unexpected runtime resolve input: %+v", runtime.lastResolved) } } diff --git a/internal/tui/core/app/update_test.go b/internal/tui/core/app/update_test.go index 68b37719..dfc2e776 100644 --- a/internal/tui/core/app/update_test.go +++ b/internal/tui/core/app/update_test.go @@ -14,6 +14,7 @@ import ( "neo-code/internal/memo" providertypes "neo-code/internal/provider/types" agentruntime "neo-code/internal/runtime" + approvalflow "neo-code/internal/runtime/approval" agentsession "neo-code/internal/session" "neo-code/internal/tools" tuibootstrap "neo-code/internal/tui/bootstrap" @@ -199,13 +200,13 @@ func TestAppUpdateBasic(t *testing.T) { } func TestParsePermissionShortcutFromKeyInput(t *testing.T) { - if decision, ok := parsePermissionShortcut("y"); !ok || decision != agentruntime.PermissionResolutionAllowOnce { + if decision, ok := parsePermissionShortcut("y"); !ok || decision != approvalflow.DecisionAllowOnce { t.Fatalf("expected allow_once, got %v (ok=%v)", decision, ok) } - if decision, ok := parsePermissionShortcut("a"); !ok || decision != agentruntime.PermissionResolutionAllowSession { + if decision, ok := parsePermissionShortcut("a"); !ok || decision != approvalflow.DecisionAllowSession { t.Fatalf("expected allow_session, got %v (ok=%v)", decision, ok) } - if decision, ok := parsePermissionShortcut("n"); !ok || decision != agentruntime.PermissionResolutionReject { + if decision, ok := parsePermissionShortcut("n"); !ok || decision != approvalflow.DecisionReject { t.Fatalf("expected reject, got %v (ok=%v)", decision, ok) } if _, ok := parsePermissionShortcut("x"); ok { @@ -280,7 +281,7 @@ func TestUpdatePermissionResolveFlow(t *testing.T) { if len(runtime.resolveCalls) != 1 || runtime.resolveCalls[0].RequestID != "perm-3" { t.Fatalf("expected ResolvePermission to be called") } - if runtime.resolveCalls[0].Decision != agentruntime.PermissionResolutionAllowOnce { + if runtime.resolveCalls[0].Decision != approvalflow.DecisionAllowOnce { t.Fatalf("unexpected decision forwarded: %s", runtime.resolveCalls[0].Decision) } @@ -303,7 +304,7 @@ func TestUpdatePermissionResolvedError(t *testing.T) { model, _ := app.Update(permissionResolutionFinishedMsg{ RequestID: "perm-4", - Decision: agentruntime.PermissionResolutionAllowOnce, + Decision: approvalflow.DecisionAllowOnce, Err: errors.New("boom"), }) app = model.(App) @@ -318,7 +319,7 @@ func TestUpdatePermissionResolvedError(t *testing.T) { func TestRunResolvePermissionCommand(t *testing.T) { runtime := newStubRuntime() - cmd := runResolvePermission(runtime, "perm-5", agentruntime.PermissionResolutionAllowSession) + cmd := runResolvePermission(runtime, "perm-5", approvalflow.DecisionAllowSession) if cmd == nil { t.Fatalf("expected command") } @@ -327,7 +328,7 @@ func TestRunResolvePermissionCommand(t *testing.T) { if !ok { t.Fatalf("expected permissionResolutionFinishedMsg, got %T", msg) } - if resolved.RequestID != "perm-5" || resolved.Decision != agentruntime.PermissionResolutionAllowSession { + if resolved.RequestID != "perm-5" || resolved.Decision != approvalflow.DecisionAllowSession { t.Fatalf("unexpected resolved msg: %#v", resolved) } if len(runtime.resolveCalls) != 1 { @@ -358,7 +359,7 @@ func TestUpdatePermissionResolutionFinishedMsgIgnoresMismatch(t *testing.T) { } model, cmd := app.Update(permissionResolutionFinishedMsg{ RequestID: "perm-8", - Decision: agentruntime.PermissionResolutionAllowOnce, + Decision: approvalflow.DecisionAllowOnce, }) if model == nil { t.Fatalf("expected model") @@ -397,7 +398,7 @@ func TestUpdatePermissionRejectFlow(t *testing.T) { msg := cmd() next, _ := app.Update(msg) app = next.(App) - if len(runtime.resolveCalls) != 1 || runtime.resolveCalls[0].Decision != agentruntime.PermissionResolutionReject { + if len(runtime.resolveCalls) != 1 || runtime.resolveCalls[0].Decision != approvalflow.DecisionReject { t.Fatalf("expected reject decision to be submitted") } if app.state.StatusText != statusPermissionSubmitted { @@ -1039,7 +1040,7 @@ func TestRuntimeEventProviderRetryHandler(t *testing.T) { func TestRuntimeEventCompactDoneHandler(t *testing.T) { app, _ := newTestApp(t) - payload := agentruntime.CompactDonePayload{TriggerMode: "auto", SavedRatio: 0.5, BeforeChars: 10, AfterChars: 5, TranscriptPath: "path"} + payload := agentruntime.CompactResult{TriggerMode: "auto", SavedRatio: 0.5, BeforeChars: 10, AfterChars: 5, TranscriptPath: "path"} handled := runtimeEventCompactDoneHandler(&app, agentruntime.RuntimeEvent{Payload: payload}) if !handled { t.Fatalf("expected true") diff --git a/internal/tui/services/services_test.go b/internal/tui/services/services_test.go index 46ff30e4..39afd6f3 100644 --- a/internal/tui/services/services_test.go +++ b/internal/tui/services/services_test.go @@ -13,6 +13,7 @@ import ( configstate "neo-code/internal/config/state" providertypes "neo-code/internal/provider/types" agentruntime "neo-code/internal/runtime" + approvalflow "neo-code/internal/runtime/approval" ) type stubRunner struct { @@ -120,7 +121,7 @@ func TestRunResolvePermissionCmd(t *testing.T) { resolver := &stubPermissionResolver{err: errors.New("permission failed")} input := agentruntime.PermissionResolutionInput{ RequestID: "perm-1", - Decision: agentruntime.PermissionResolutionAllowSession, + Decision: approvalflow.DecisionAllowSession, } msg := RunResolvePermissionCmd( resolver, @@ -140,13 +141,13 @@ func TestRunResolvePermissionCmd(t *testing.T) { if !ok { t.Fatalf("expected wrapped permission result message, got %T %#v", msg, msg) } - if got.Input.RequestID != "perm-1" || got.Input.Decision != agentruntime.PermissionResolutionAllowSession { + if got.Input.RequestID != "perm-1" || got.Input.Decision != approvalflow.DecisionAllowSession { t.Fatalf("unexpected permission input forwarded: %+v", got.Input) } if got.Err == nil || got.Err.Error() != "permission failed" { t.Fatalf("expected forwarded permission error, got %#v", got.Err) } - if resolver.lastInput.RequestID != "perm-1" || resolver.lastInput.Decision != agentruntime.PermissionResolutionAllowSession { + if resolver.lastInput.RequestID != "perm-1" || resolver.lastInput.Decision != approvalflow.DecisionAllowSession { t.Fatalf("unexpected resolver input: %+v", resolver.lastInput) } if !resolver.hasDeadline { From 7345f7630a1fca8d4cd93d4c4557b4c1523d4284 Mon Sep 17 00:00:00 2001 From: xgopilot Date: Mon, 13 Apr 2026 05:53:13 +0000 Subject: [PATCH 2/5] fix(runtime): drain streaming events after cancellation Keep the stream consumer alive until the producer closes the channel so provider sends cannot block when context is canceled. Also add a regression test that reproduces cancellation with post-cancel event bursts. Generated with [codeagent](https://github.com/qbox/codeagent) Co-authored-by: phantom5099 <245659304+phantom5099@users.noreply.github.com> --- internal/runtime/provider_stream.go | 38 ++++++++----------- .../runtime_remaining_branches_test.go | 33 ++++++++++++++++ 2 files changed, 48 insertions(+), 23 deletions(-) diff --git a/internal/runtime/provider_stream.go b/internal/runtime/provider_stream.go index 97dd5992..3cbb507c 100644 --- a/internal/runtime/provider_stream.go +++ b/internal/runtime/provider_stream.go @@ -34,29 +34,21 @@ func generateStreamingMessage( streamDone <- outcome }() - for { - select { - case event, ok := <-streamEvents: - if !ok { - return - } - if err := streaming.HandleEvent(event, acc, streaming.Hooks{ - OnTextDelta: hooks.OnTextDelta, - OnToolCallStart: hooks.OnToolCallStart, - OnMessageDone: func(payload providertypes.MessageDonePayload) { - if payload.Usage != nil { - outcome.inputTokens = payload.Usage.InputTokens - outcome.outputTokens = payload.Usage.OutputTokens - } - if hooks.OnMessageDone != nil { - hooks.OnMessageDone(payload) - } - }, - }); err != nil && outcome.err == nil { - outcome.err = err - } - case <-ctx.Done(): - return + userOnMessageDone := hooks.OnMessageDone + hooksCopy := hooks + hooksCopy.OnMessageDone = func(payload providertypes.MessageDonePayload) { + if payload.Usage != nil { + outcome.inputTokens = payload.Usage.InputTokens + outcome.outputTokens = payload.Usage.OutputTokens + } + if userOnMessageDone != nil { + userOnMessageDone(payload) + } + } + + for event := range streamEvents { + if err := streaming.HandleEvent(event, acc, hooksCopy); err != nil && outcome.err == nil { + outcome.err = err } } }() diff --git a/internal/runtime/runtime_remaining_branches_test.go b/internal/runtime/runtime_remaining_branches_test.go index d618a5f1..01dadba4 100644 --- a/internal/runtime/runtime_remaining_branches_test.go +++ b/internal/runtime/runtime_remaining_branches_test.go @@ -3,6 +3,7 @@ package runtime import ( "context" "errors" + "fmt" "os" "sync" "testing" @@ -138,6 +139,38 @@ func TestGenerateStreamingMessageHooksAndContextDoneBranches(t *testing.T) { } } +func TestGenerateStreamingMessageDrainEventsAfterContextCanceled(t *testing.T) { + t.Parallel() + + ctx, cancel := context.WithCancel(context.Background()) + providerIgnoreCancel := &scriptedProvider{chatFn: func(ctx context.Context, req providertypes.GenerateRequest, events chan<- providertypes.StreamEvent) error { + <-ctx.Done() + for i := 0; i < 64; i++ { + events <- providertypes.NewTextDeltaStreamEvent(fmt.Sprintf("delta-%d", i)) + } + return context.Canceled + }} + + go func() { + time.Sleep(10 * time.Millisecond) + cancel() + }() + + done := make(chan streamGenerateResult, 1) + go func() { + done <- generateStreamingMessage(ctx, providerIgnoreCancel, providertypes.GenerateRequest{}, streaming.Hooks{}) + }() + + select { + case outcome := <-done: + if !errors.Is(outcome.err, context.Canceled) { + t.Fatalf("expected context canceled, got %v", outcome.err) + } + case <-time.After(2 * time.Second): + t.Fatal("generateStreamingMessage blocked after context cancellation") + } +} + func TestPrepareTurnSnapshotErrorBranches(t *testing.T) { t.Parallel() From 27ec07a79d39591d7cef58ac0354b8e865fe920f Mon Sep 17 00:00:00 2001 From: xgopilot Date: Mon, 13 Apr 2026 06:07:01 +0000 Subject: [PATCH 3/5] fix(runtime): preserve memo extractor hook after refactor Generated with [codeagent](https://github.com/qbox/codeagent) Co-authored-by: phantom5099 <245659304+phantom5099@users.noreply.github.com> --- internal/runtime/memo.go | 17 +++++ internal/runtime/run.go | 1 + internal/runtime/runtime.go | 14 ++++ .../runtime/runtime_internal_helpers_test.go | 66 +++++++++++++++++++ 4 files changed, 98 insertions(+) create mode 100644 internal/runtime/memo.go diff --git a/internal/runtime/memo.go b/internal/runtime/memo.go new file mode 100644 index 00000000..6404b631 --- /dev/null +++ b/internal/runtime/memo.go @@ -0,0 +1,17 @@ +package runtime + +import ( + "context" + + providertypes "neo-code/internal/provider/types" +) + +// triggerMemoExtraction 在 Run 结束后异步触发记忆提取，避免阻塞主闭环。 +func (s *Service) triggerMemoExtraction(messages []providertypes.Message) { + if s == nil || s.memoExtractor == nil || len(messages) == 0 { + return + } + + cloned := append([]providertypes.Message(nil), messages...) + go s.memoExtractor.ExtractAndStore(context.Background(), cloned) +} diff --git a/internal/runtime/run.go b/internal/runtime/run.go index 0821fb9c..97926939 100644 --- a/internal/runtime/run.go +++ b/internal/runtime/run.go @@ -90,6 +90,7 @@ func (s *Service) Run(ctx context.Context, input UserInput) error { if len(turnResult.assistant.ToolCalls) == 0 { s.emit(ctx, EventAgentDone, state.runID, state.session.ID, turnResult.assistant) + s.triggerMemoExtraction(state.session.Messages) return nil } if err := s.executeAssistantToolCalls(ctx, &state, snapshot, turnResult.assistant); err != nil { diff --git a/internal/runtime/runtime.go b/internal/runtime/runtime.go index d5738da0..2027571a 100644 --- a/internal/runtime/runtime.go +++ b/internal/runtime/runtime.go @@ -11,6 +11,7 @@ import ( agentcontext "neo-code/internal/context" contextcompact "neo-code/internal/context/compact" "neo-code/internal/provider" + providertypes "neo-code/internal/provider/types" "neo-code/internal/runtime/approval" agentsession "neo-code/internal/session" "neo-code/internal/tools" @@ -47,6 +48,13 @@ type ProviderFactory interface { Build(ctx context.Context, cfg provider.RuntimeConfig) (provider.Provider, error) } +// MemoExtractor 定义 runtime 层调用记忆提取的最小能力。 +// 通过接口注入避免 runtime 直接依赖 memo 子系统实现细节。 +type MemoExtractor interface { + // ExtractAndStore 从对话消息中提取并落盘记忆，失败由实现方自行处理。 + ExtractAndStore(ctx context.Context, messages []providertypes.Message) +} + // Service 是 runtime 的默认实现，负责组织一次完整的 agent 运行闭环。 type Service struct { configManager *config.Manager @@ -56,6 +64,7 @@ type Service struct { contextBuilder agentcontext.Builder compactRunner contextcompact.Runner approvalBroker *approval.Broker + memoExtractor MemoExtractor events chan RuntimeEvent sessionMu sync.Mutex @@ -96,6 +105,11 @@ func NewWithFactory( } } +// SetMemoExtractor 设置可选记忆提取钩子，由 Run 在结束时异步触发。 +func (s *Service) SetMemoExtractor(extractor MemoExtractor) { + s.memoExtractor = extractor +} + // CancelActiveRun 尝试取消当前正在执行的 Run。 func (s *Service) CancelActiveRun() bool { s.runMu.Lock() diff --git a/internal/runtime/runtime_internal_helpers_test.go b/internal/runtime/runtime_internal_helpers_test.go index c9d72647..059bf554 100644 --- a/internal/runtime/runtime_internal_helpers_test.go +++ b/internal/runtime/runtime_internal_helpers_test.go @@ -3,6 +3,7 @@ package runtime import ( "context" "errors" + "sync" "testing" "time" @@ -14,6 +15,27 @@ import ( "neo-code/internal/tools" ) +type stubMemoExtractor struct { + mu sync.Mutex + calls int + lastMsgs []providertypes.Message + doneCh chan struct{} +} + +func (s *stubMemoExtractor) ExtractAndStore(_ context.Context, messages []providertypes.Message) { + s.mu.Lock() + s.calls++ + s.lastMsgs = append([]providertypes.Message(nil), messages...) + doneCh := s.doneCh + s.mu.Unlock() + if doneCh != nil { + select { + case doneCh <- struct{}{}: + default: + } + } +} + func TestRunStateNilReceiverNoops(t *testing.T) { t.Parallel() @@ -228,6 +250,50 @@ func TestExecuteAssistantToolCallsCanceledSaveStillEmitsResultWhenExecErr(t *tes assertEventSequence(t, events, []EventType{EventToolStart, EventToolResult}) } +func TestSetMemoExtractorAndRunTriggersExtraction(t *testing.T) { + t.Parallel() + + store := newMemoryStore() + providerStub := &scriptedProvider{ + streams: [][]providertypes.StreamEvent{ + { + providertypes.NewTextDeltaStreamEvent("memo ready"), + providertypes.NewMessageDoneStreamEvent("", nil), + }, + }, + } + factory := &scriptedProviderFactory{provider: providerStub} + toolManager := &stubToolManager{} + service := NewWithFactory( + newRuntimeConfigManagerWithProviderEnvs(t, nil), + toolManager, + store, + factory, + &stubContextBuilder{}, + ) + extractor := &stubMemoExtractor{doneCh: make(chan struct{}, 1)} + service.SetMemoExtractor(extractor) + + if err := service.Run(context.Background(), UserInput{RunID: "run-memo-extract", Content: "hello"}); err != nil { + t.Fatalf("Run() error = %v", err) + } + + select { + case <-extractor.doneCh: + case <-time.After(2 * time.Second): + t.Fatalf("memo extractor was not triggered") + } + + extractor.mu.Lock() + defer extractor.mu.Unlock() + if extractor.calls != 1 { + t.Fatalf("expected memo extractor to be called once, got %d", extractor.calls) + } + if len(extractor.lastMsgs) < 2 { + t.Fatalf("expected user+assistant messages, got %d", len(extractor.lastMsgs)) + } +} + func newRuntimeSession(id string) agentsession.Session { session := agentsession.New("runtime test") session.ID = id From a80d7b85081f2f7573fe1f5a422f3e39fc5180b5 Mon Sep 17 00:00:00 2001 From: phantom5099 <1011668688@qq.com> Date: Mon, 13 Apr 2026 16:29:20 +0800 Subject: [PATCH 4/5] =?UTF-8?q?refactor(runtime):=E4=BC=98=E5=8C=96?= =?UTF-8?q?=E5=90=8C=E4=B8=80=20Session=20=E5=B9=B6=E5=8F=91=20Run=20?= =?UTF-8?q?=E4=BC=9A=E4=B8=A2=E6=B6=88=E6=81=AF/=E8=A6=86=E7=9B=96?= =?UTF-8?q?=E7=8A=B6=E6=80=81?= MIME-Version: 1.0 Content-Type: text/plain; charset=UTF-8 Content-Transfer-Encoding: 8bit --- internal/runtime/run.go | 24 ++- internal/runtime/runtime.go | 50 +++--- internal/runtime/runtime_gap_coverage_test.go | 25 ++- internal/runtime/runtime_test.go | 150 ++++++++++++++++++ 4 files changed, 222 insertions(+), 27 deletions(-) diff --git a/internal/runtime/run.go b/internal/runtime/run.go index 97926939..243ef2e1 100644 --- a/internal/runtime/run.go +++ b/internal/runtime/run.go @@ -18,8 +18,8 @@ import ( ) // Run 执行一次完整的 ReAct 闭环：保存用户输入、驱动模型、执行工具并发出事件。 -// 会话创建/加载在锁外完成，仅会话级操作（ReAct 循环）受会话锁保护， -// 不同会话可并行执行。 +// 已有会话会先加锁再加载/更新，确保同一会话并发 Run 不会出现状态覆盖； +// 新会话在创建后再绑定会话锁，不同会话可并行执行。 func (s *Service) Run(ctx context.Context, input UserInput) error { if strings.TrimSpace(input.Content) == "" { return errors.New("runtime: input content is empty") @@ -34,14 +34,28 @@ func (s *Service) Run(ctx context.Context, input UserInput) error { ctx = runCtx initialCfg := s.configManager.Get() + sessionID := strings.TrimSpace(input.SessionID) + releaseSessionLock := func() {} + defer func() { + releaseSessionLock() + }() + + if sessionID != "" { + sessionMu := s.acquireSessionLock(sessionID) + sessionMu.Lock() + releaseSessionLock = sessionMu.Unlock + } + session, err := s.loadOrCreateSession(ctx, input.SessionID, input.Content, initialCfg.Workdir, input.Workdir) if err != nil { return s.handleRunError(ctx, input.RunID, input.SessionID, err) } - sessionMu := s.acquireSessionLock(session.ID) - sessionMu.Lock() - defer sessionMu.Unlock() + if sessionID == "" { + sessionMu := s.acquireSessionLock(session.ID) + sessionMu.Lock() + releaseSessionLock = sessionMu.Unlock + } state := newRunState(input.RunID, session) if err := s.appendUserMessageAndSave(ctx, &state, input.Content); err != nil { diff --git a/internal/runtime/runtime.go b/internal/runtime/runtime.go index 2027571a..4de75f2a 100644 --- a/internal/runtime/runtime.go +++ b/internal/runtime/runtime.go @@ -66,13 +66,13 @@ type Service struct { approvalBroker *approval.Broker memoExtractor MemoExtractor - events chan RuntimeEvent - sessionMu sync.Mutex - sessionLocks map[string]*sync.Mutex - runMu sync.Mutex - activeRunToken uint64 - nextRunToken uint64 - activeRunCancel context.CancelFunc + events chan RuntimeEvent + sessionMu sync.Mutex + sessionLocks map[string]*sync.Mutex + runMu sync.Mutex + activeRunToken uint64 + nextRunToken uint64 + activeRunCancels map[uint64]context.CancelFunc } // NewWithFactory 使用注入依赖构建默认 runtime Service。 @@ -94,14 +94,15 @@ func NewWithFactory( } return &Service{ - configManager: configManager, - sessionStore: sessionStore, - toolManager: toolManager, - providerFactory: providerFactory, - contextBuilder: contextBuilder, - approvalBroker: approval.NewBroker(), - events: make(chan RuntimeEvent, 128), - sessionLocks: make(map[string]*sync.Mutex), + configManager: configManager, + sessionStore: sessionStore, + toolManager: toolManager, + providerFactory: providerFactory, + contextBuilder: contextBuilder, + approvalBroker: approval.NewBroker(), + events: make(chan RuntimeEvent, 128), + sessionLocks: make(map[string]*sync.Mutex), + activeRunCancels: make(map[uint64]context.CancelFunc), } } @@ -110,10 +111,14 @@ func (s *Service) SetMemoExtractor(extractor MemoExtractor) { s.memoExtractor = extractor } -// CancelActiveRun 尝试取消当前正在执行的 Run。 +// CancelActiveRun 尝试取消最近一次仍在执行的 Run。 func (s *Service) CancelActiveRun() bool { s.runMu.Lock() - cancel := s.activeRunCancel + if s.activeRunToken == 0 { + s.runMu.Unlock() + return false + } + cancel := s.activeRunCancels[s.activeRunToken] s.runMu.Unlock() if cancel == nil { return false @@ -215,21 +220,26 @@ func (s *Service) startRun(cancel context.CancelFunc) uint64 { s.nextRunToken++ token := s.nextRunToken s.activeRunToken = token - s.activeRunCancel = cancel + s.activeRunCancels[token] = cancel return token } -// finishRun 在运行结束时释放当前激活的取消句柄。 +// finishRun 在运行结束时释放指定运行的取消句柄，并回退到最新活跃运行。 func (s *Service) finishRun(token uint64) { s.runMu.Lock() defer s.runMu.Unlock() + delete(s.activeRunCancels, token) if s.activeRunToken != token { return } s.activeRunToken = 0 - s.activeRunCancel = nil + for activeToken := range s.activeRunCancels { + if activeToken > s.activeRunToken { + s.activeRunToken = activeToken + } + } } // acquireSessionLock 获取指定会话的互斥锁，确保同一会话的 Run/Compact 串行执行。 diff --git a/internal/runtime/runtime_gap_coverage_test.go b/internal/runtime/runtime_gap_coverage_test.go index 6eb03dcc..2de5841a 100644 --- a/internal/runtime/runtime_gap_coverage_test.go +++ b/internal/runtime/runtime_gap_coverage_test.go @@ -218,12 +218,33 @@ func TestRuntimeLoadOrCreateAndEmitBranches(t *testing.T) { } <-done - token1 := service.startRun(func() {}) - token2 := service.startRun(func() {}) + canceled := make([]string, 0, 2) + token1 := service.startRun(func() { canceled = append(canceled, "run-1") }) + token2 := service.startRun(func() { canceled = append(canceled, "run-2") }) service.finishRun(token1) if service.activeRunToken != token2 { t.Fatalf("expected finishRun with stale token to keep active run") } + + if !service.CancelActiveRun() { + t.Fatalf("expected cancel latest active run") + } + if len(canceled) != 1 || canceled[0] != "run-2" { + t.Fatalf("expected latest run canceled first, got %v", canceled) + } + + service.finishRun(token2) + if !service.CancelActiveRun() { + t.Fatalf("expected previous active run to become cancel target") + } + if len(canceled) != 2 || canceled[1] != "run-1" { + t.Fatalf("expected second cancel to target run-1, got %v", canceled) + } + + service.finishRun(token1) + if service.CancelActiveRun() { + t.Fatalf("expected no active run after all runs finished") + } } func TestPermissionHelperAndAwaitBranches(t *testing.T) { diff --git a/internal/runtime/runtime_test.go b/internal/runtime/runtime_test.go index b0667462..f10d7b78 100644 --- a/internal/runtime/runtime_test.go +++ b/internal/runtime/runtime_test.go @@ -7,6 +7,7 @@ import ( "os" "path/filepath" "strings" + "sync" "testing" "time" @@ -92,6 +93,81 @@ func (s *memoryStore) ListSummaries(ctx context.Context) ([]agentsession.Summary return summaries, nil } +// blockingLoadStore 用于并发测试：首次 Load 阻塞，以验证同 session 的锁时序。 +type blockingLoadStore struct { + mu sync.Mutex + sessions map[string]agentsession.Session + loadCalls int + loadEntered chan struct{} + unblockFirst chan struct{} + loadEnteredSet bool +} + +func newBlockingLoadStore() *blockingLoadStore { + return &blockingLoadStore{ + sessions: map[string]agentsession.Session{}, + loadEntered: make(chan struct{}), + unblockFirst: make(chan struct{}), + } +} + +func (s *blockingLoadStore) Save(ctx context.Context, session *agentsession.Session) error { + if err := ctx.Err(); err != nil { + return err + } + if session == nil { + return errors.New("nil session") + } + s.mu.Lock() + s.sessions[session.ID] = cloneSession(*session) + s.mu.Unlock() + return nil +} + +func (s *blockingLoadStore) Load(ctx context.Context, id string) (agentsession.Session, error) { + if err := ctx.Err(); err != nil { + return agentsession.Session{}, err + } + s.mu.Lock() + s.loadCalls++ + callIndex := s.loadCalls + if callIndex == 1 && !s.loadEnteredSet { + s.loadEnteredSet = true + close(s.loadEntered) + } + s.mu.Unlock() + + if callIndex == 1 { + <-s.unblockFirst + } + + s.mu.Lock() + session, ok := s.sessions[id] + s.mu.Unlock() + if !ok { + return agentsession.Session{}, errors.New("not found") + } + return cloneSession(session), nil +} + +func (s *blockingLoadStore) ListSummaries(ctx context.Context) ([]agentsession.Summary, error) { + if err := ctx.Err(); err != nil { + return nil, err + } + s.mu.Lock() + defer s.mu.Unlock() + summaries := make([]agentsession.Summary, 0, len(s.sessions)) + for _, session := range s.sessions { + summaries = append(summaries, agentsession.Summary{ + ID: session.ID, + Title: session.Title, + CreatedAt: session.CreatedAt, + UpdatedAt: session.UpdatedAt, + }) + } + return summaries, nil +} + type scriptedProvider struct { name string streams [][]providertypes.StreamEvent @@ -1664,6 +1740,80 @@ func TestServiceCancelActiveRun(t *testing.T) { assertEventsRunID(t, events, input.RunID) } +func TestServiceRunSameSessionConcurrentNoMessageLoss(t *testing.T) { + manager := newRuntimeConfigManager(t) + registry := tools.NewRegistry() + registry.Register(&stubTool{name: "filesystem_read_file", content: "default"}) + + store := newBlockingLoadStore() + session := agentsession.New("same-session") + session.ID = "session-same-concurrent" + store.sessions[session.ID] = cloneSession(session) + + scripted := &scriptedProvider{ + name: "same-session-concurrent-provider", + chatFn: func(ctx context.Context, req providertypes.GenerateRequest, events chan<- providertypes.StreamEvent) error { + events <- providertypes.NewTextDeltaStreamEvent("done") + events <- providertypes.NewMessageDoneStreamEvent("stop", nil) + return nil + }, + } + + service := NewWithFactory(manager, registry, store, &scriptedProviderFactory{provider: scripted}, nil) + errCh1 := make(chan error, 1) + errCh2 := make(chan error, 1) + go func() { + errCh1 <- service.Run(context.Background(), UserInput{SessionID: session.ID, RunID: "run-same-1", Content: "hello-1"}) + }() + + select { + case <-store.loadEntered: + case <-time.After(2 * time.Second): + t.Fatalf("timed out waiting first load entry") + } + + go func() { + errCh2 <- service.Run(context.Background(), UserInput{SessionID: session.ID, RunID: "run-same-2", Content: "hello-2"}) + }() + + time.Sleep(120 * time.Millisecond) + store.mu.Lock() + loadCallsBeforeRelease := store.loadCalls + store.mu.Unlock() + if loadCallsBeforeRelease != 1 { + t.Fatalf("expected second run not to load before first lock release, got load calls = %d", loadCallsBeforeRelease) + } + + close(store.unblockFirst) + if err := <-errCh1; err != nil { + t.Fatalf("first run error = %v", err) + } + if err := <-errCh2; err != nil { + t.Fatalf("second run error = %v", err) + } + + store.mu.Lock() + finalSession := cloneSession(store.sessions[session.ID]) + store.mu.Unlock() + + if len(finalSession.Messages) != 4 { + t.Fatalf("expected 4 messages from two runs, got %+v", finalSession.Messages) + } + userCount := 0 + assistantCount := 0 + for _, message := range finalSession.Messages { + switch message.Role { + case providertypes.RoleUser: + userCount++ + case providertypes.RoleAssistant: + assistantCount++ + } + } + if userCount != 2 || assistantCount != 2 { + t.Fatalf("expected 2 user + 2 assistant messages, got users=%d assistants=%d messages=%+v", userCount, assistantCount, finalSession.Messages) + } +} + func TestServiceRunCanceledByProvider(t *testing.T) { manager := newRuntimeConfigManager(t) store := newMemoryStore() From 5fe066e082a62c8340f2f5efc7027c5365b862d3 Mon Sep 17 00:00:00 2001 From: xgopilot Date: Mon, 13 Apr 2026 08:37:08 +0000 Subject: [PATCH 5/5] fix(runtime): recycle session locks and harden nil maps Generated with [codeagent](https://github.com/qbox/codeagent) Co-authored-by: phantom5099 <245659304+phantom5099@users.noreply.github.com> --- internal/runtime/compact.go | 7 ++- internal/runtime/run.go | 14 +++-- internal/runtime/runtime.go | 52 +++++++++++++++---- internal/runtime/runtime_gap_coverage_test.go | 39 +++++++++++--- 4 files changed, 89 insertions(+), 23 deletions(-) diff --git a/internal/runtime/compact.go b/internal/runtime/compact.go index b7a1227d..0d3e7843 100644 --- a/internal/runtime/compact.go +++ b/internal/runtime/compact.go @@ -66,9 +66,12 @@ func (s *Service) Compact(ctx context.Context, input CompactInput) (CompactResul return CompactResult{}, errors.New("runtime: compact session_id is empty") } - sessionMu := s.acquireSessionLock(input.SessionID) + sessionMu, releaseLockRef := s.acquireSessionLock(input.SessionID) sessionMu.Lock() - defer sessionMu.Unlock() + defer func() { + sessionMu.Unlock() + releaseLockRef() + }() cfg := s.configManager.Get() session, err := s.sessionStore.Load(ctx, input.SessionID) diff --git a/internal/runtime/run.go b/internal/runtime/run.go index 243ef2e1..c1872192 100644 --- a/internal/runtime/run.go +++ b/internal/runtime/run.go @@ -41,9 +41,12 @@ func (s *Service) Run(ctx context.Context, input UserInput) error { }() if sessionID != "" { - sessionMu := s.acquireSessionLock(sessionID) + sessionMu, releaseLockRef := s.acquireSessionLock(sessionID) sessionMu.Lock() - releaseSessionLock = sessionMu.Unlock + releaseSessionLock = func() { + sessionMu.Unlock() + releaseLockRef() + } } session, err := s.loadOrCreateSession(ctx, input.SessionID, input.Content, initialCfg.Workdir, input.Workdir) @@ -52,9 +55,12 @@ func (s *Service) Run(ctx context.Context, input UserInput) error { } if sessionID == "" { - sessionMu := s.acquireSessionLock(session.ID) + sessionMu, releaseLockRef := s.acquireSessionLock(session.ID) sessionMu.Lock() - releaseSessionLock = sessionMu.Unlock + releaseSessionLock = func() { + sessionMu.Unlock() + releaseLockRef() + } } state := newRunState(input.RunID, session) diff --git a/internal/runtime/runtime.go b/internal/runtime/runtime.go index 4de75f2a..a49f25a7 100644 --- a/internal/runtime/runtime.go +++ b/internal/runtime/runtime.go @@ -68,13 +68,19 @@ type Service struct { events chan RuntimeEvent sessionMu sync.Mutex - sessionLocks map[string]*sync.Mutex + sessionLocks map[string]*sessionLockEntry runMu sync.Mutex activeRunToken uint64 nextRunToken uint64 activeRunCancels map[uint64]context.CancelFunc } +// sessionLockEntry 维护单个会话锁及其当前引用计数，用于在无引用时回收 map 项。 +type sessionLockEntry struct { + mu sync.Mutex + refs int +} + // NewWithFactory 使用注入依赖构建默认 runtime Service。 func NewWithFactory( configManager *config.Manager, @@ -101,7 +107,7 @@ func NewWithFactory( contextBuilder: contextBuilder, approvalBroker: approval.NewBroker(), events: make(chan RuntimeEvent, 128), - sessionLocks: make(map[string]*sync.Mutex), + sessionLocks: make(map[string]*sessionLockEntry), activeRunCancels: make(map[uint64]context.CancelFunc), } } @@ -216,6 +222,9 @@ func (s *Service) emit(ctx context.Context, kind EventType, runID string, sessio func (s *Service) startRun(cancel context.CancelFunc) uint64 { s.runMu.Lock() defer s.runMu.Unlock() + if s.activeRunCancels == nil { + s.activeRunCancels = make(map[uint64]context.CancelFunc) + } s.nextRunToken++ token := s.nextRunToken @@ -242,18 +251,41 @@ func (s *Service) finishRun(token uint64) { } } -// acquireSessionLock 获取指定会话的互斥锁，确保同一会话的 Run/Compact 串行执行。 -// 不同会话的锁互不干扰，支持多会话并行。 -func (s *Service) acquireSessionLock(sessionID string) *sync.Mutex { +// acquireSessionLock 获取指定会话锁并返回释放引用的函数。 +// 调用方在完成会话级串行操作后，必须调用 release 以允许锁条目回收。 +func (s *Service) acquireSessionLock(sessionID string) (*sync.Mutex, func()) { s.sessionMu.Lock() - defer s.sessionMu.Unlock() + if s.sessionLocks == nil { + s.sessionLocks = make(map[string]*sessionLockEntry) + } - mu, ok := s.sessionLocks[sessionID] + entry, ok := s.sessionLocks[sessionID] if !ok { - mu = &sync.Mutex{} - s.sessionLocks[sessionID] = mu + entry = &sessionLockEntry{} + s.sessionLocks[sessionID] = entry + } + entry.refs++ + s.sessionMu.Unlock() + + released := false + release := func() { + s.sessionMu.Lock() + defer s.sessionMu.Unlock() + if released { + return + } + released = true + + current, exists := s.sessionLocks[sessionID] + if !exists || current != entry { + return + } + current.refs-- + if current.refs <= 0 { + delete(s.sessionLocks, sessionID) + } } - return mu + return &entry.mu, release } func resolveWorkdirForSession(defaultWorkdir string, currentWorkdir string, requestedWorkdir string) (string, error) { diff --git a/internal/runtime/runtime_gap_coverage_test.go b/internal/runtime/runtime_gap_coverage_test.go index 2de5841a..692a0f88 100644 --- a/internal/runtime/runtime_gap_coverage_test.go +++ b/internal/runtime/runtime_gap_coverage_test.go @@ -3,7 +3,6 @@ package runtime import ( "context" "errors" - "sync" "testing" "time" @@ -24,7 +23,7 @@ func TestCompactValidationAndLoadErrorBranches(t *testing.T) { configManager: newRuntimeConfigManager(t), sessionStore: newMemoryStore(), events: make(chan RuntimeEvent, 16), - sessionLocks: make(map[string]*sync.Mutex), + sessionLocks: make(map[string]*sessionLockEntry), } ctx, cancel := context.WithCancel(context.Background()) @@ -174,7 +173,7 @@ func TestRuntimeLoadOrCreateAndEmitBranches(t *testing.T) { t.Parallel() manager := newRuntimeConfigManager(t) - service := &Service{configManager: manager, sessionStore: newMemoryStore(), events: make(chan RuntimeEvent, 1), sessionLocks: map[string]*sync.Mutex{}} + service := &Service{configManager: manager, sessionStore: newMemoryStore(), events: make(chan RuntimeEvent, 1), sessionLocks: map[string]*sessionLockEntry{}} if _, err := service.loadOrCreateSession(context.Background(), "", "title", t.TempDir(), "missing/subdir"); err == nil { t.Fatalf("expected new session workdir resolve error") @@ -234,11 +233,11 @@ func TestRuntimeLoadOrCreateAndEmitBranches(t *testing.T) { } service.finishRun(token2) - if !service.CancelActiveRun() { - t.Fatalf("expected previous active run to become cancel target") + if service.CancelActiveRun() { + t.Fatalf("expected no active run after latest run finished") } - if len(canceled) != 2 || canceled[1] != "run-1" { - t.Fatalf("expected second cancel to target run-1, got %v", canceled) + if len(canceled) != 1 || canceled[0] != "run-2" { + t.Fatalf("expected only run-2 canceled, got %v", canceled) } service.finishRun(token1) @@ -276,6 +275,32 @@ func TestPermissionHelperAndAwaitBranches(t *testing.T) { } } +func TestAcquireSessionLockReleasesAndCleansUp(t *testing.T) { + t.Parallel() + + service := &Service{} + lock1, release1 := service.acquireSessionLock("session-1") + lock2, release2 := service.acquireSessionLock("session-1") + if lock1 != lock2 { + t.Fatalf("expected same mutex instance for same session") + } + if got := service.sessionLocks["session-1"].refs; got != 2 { + t.Fatalf("expected refs=2, got %d", got) + } + + lock1.Lock() + lock1.Unlock() + release1() + if got := service.sessionLocks["session-1"].refs; got != 1 { + t.Fatalf("expected refs=1 after first release, got %d", got) + } + + release2() + if len(service.sessionLocks) != 0 { + t.Fatalf("expected session lock entry to be cleaned up") + } +} + func permissionDecisionAskError(t *testing.T) *tools.PermissionDecisionError { t.Helper()