diff --git a/arkruntime/lib/environments/integration_test.go b/arkruntime/lib/environments/integration_test.go index bead2c1..e4c012e 100644 --- a/arkruntime/lib/environments/integration_test.go +++ b/arkruntime/lib/environments/integration_test.go @@ -101,6 +101,11 @@ func (f *fakeEnvironmentWorkerAPI) SendEvent(_ context.Context, req selfhosted.S f.mu.Lock() defer f.mu.Unlock() f.sent = append(f.sent, req.Event) + f.events = append(f.events, selfhosted.Event{ + ID: "evt_idle_after_result", + Type: selfhosted.EventTypeSessionStatusIdle, + StopReason: &selfhosted.SessionStopReason{Type: selfhosted.SessionStopReasonEndTurn}, + }) return nil } @@ -124,9 +129,11 @@ func TestEnvironmentWorkerRunHandlesPolledWorkInProcess(t *testing.T) { EvaluatedPermission: selfhosted.PermissionAllow, }, { - ID: "evt_idle", - Type: selfhosted.EventTypeSessionStatusIdle, - StopReason: &selfhosted.SessionStopReason{Type: selfhosted.SessionStopReasonEndTurn}, + ID: "evt_idle", + Type: selfhosted.EventTypeSessionStatusIdle, + StopReason: &selfhosted.SessionStopReason{ + Type: selfhosted.SessionStopReasonRequiresAction, EventIDs: []string{"toolu_local"}, + }, }, }, onStop: cancel, @@ -174,14 +181,15 @@ func TestEnvironmentWorkerForwardsToolTimeout(t *testing.T) { ID: "toolu_timeout", Type: selfhosted.EventTypeAgentCustomToolUse, Name: tool.Name(), - ToolUseID: "call_timeout", Input: selfhosted.RawJSON(`{}`), SessionThreadID: "thread_timeout", }, { - ID: "evt_idle", - Type: selfhosted.EventTypeSessionStatusIdle, - StopReason: &selfhosted.SessionStopReason{Type: selfhosted.SessionStopReasonEndTurn}, + ID: "evt_idle", + Type: selfhosted.EventTypeSessionStatusIdle, + StopReason: &selfhosted.SessionStopReason{ + Type: selfhosted.SessionStopReasonRequiresAction, EventIDs: []string{"toolu_timeout"}, + }, }, }, } @@ -226,9 +234,11 @@ func TestEnvironmentWorkerForwardsToolTimeoutToDefaultTools(t *testing.T) { EvaluatedPermission: selfhosted.PermissionAllow, }, { - ID: "evt_idle", - Type: selfhosted.EventTypeSessionStatusIdle, - StopReason: &selfhosted.SessionStopReason{Type: selfhosted.SessionStopReasonEndTurn}, + ID: "evt_idle", + Type: selfhosted.EventTypeSessionStatusIdle, + StopReason: &selfhosted.SessionStopReason{ + Type: selfhosted.SessionStopReasonRequiresAction, EventIDs: []string{"toolu_default_timeout"}, + }, }, }, } diff --git a/arkruntime/selfhosted/session_tool_runner.go b/arkruntime/selfhosted/session_tool_runner.go index 2eb9e96..9c9fc9e 100644 --- a/arkruntime/selfhosted/session_tool_runner.go +++ b/arkruntime/selfhosted/session_tool_runner.go @@ -50,6 +50,11 @@ type ToolResultStore interface { MarkSent(callID string) error } +// ToolResultStoreDiscarder 可选地从持久化 store 中删除不再属于当前 session 阻塞集合的恢复结果。 +type ToolResultStoreDiscarder interface { + Discard(callID string) error +} + // SessionToolRunnerOptions 配置单 session 的 tool event loop。 type SessionToolRunnerOptions struct { WorkID string @@ -180,15 +185,19 @@ func (r *SessionToolRunner) runLoop() error { return errors.New("session tool runner tools must not be nil") } state := &toolRunnerState{ - runner: r, - page: r.opts.EventPage, - processed: map[string]bool{}, - seen: map[string]bool{}, - answered: map[string]bool{}, - pendingResults: map[string]Event{}, - pendingAsk: map[string]Event{}, - confirmations: map[string]Event{}, - externalTools: map[string]Event{}, + runner: r, + page: r.opts.EventPage, + processed: map[string]bool{}, + seen: map[string]bool{}, + answered: map[string]bool{}, + pendingResults: map[string]Event{}, + recoveredResults: map[string]bool{}, + pendingAsk: map[string]Event{}, + confirmations: map[string]Event{}, + externalTools: map[string]Event{}, + sessionToolUses: map[string]bool{}, + toolUsesSinceStatus: map[string]bool{}, + blockingEventIDs: map[string]bool{}, } if r.opts.ResultStore != nil { pending, processed, err := r.opts.ResultStore.Recover() @@ -196,6 +205,9 @@ func (r *SessionToolRunner) runLoop() error { return fmt.Errorf("recover tool result store: %w", err) } state.pendingResults = pending + for id := range pending { + state.recoveredResults[id] = true + } state.processed = processed for id := range processed { state.answered[id] = true @@ -219,17 +231,22 @@ func (r *SessionToolRunner) runLoop() error { } type toolRunnerState struct { - runner *SessionToolRunner - page string - processed map[string]bool - seen map[string]bool - answered map[string]bool - pendingResults map[string]Event - pendingAsk map[string]Event - confirmations map[string]Event - externalTools map[string]Event - idleArmedAt time.Time - idleArmPending bool + runner *SessionToolRunner + page string + processed map[string]bool + seen map[string]bool + answered map[string]bool + pendingResults map[string]Event + recoveredResults map[string]bool + pendingAsk map[string]Event + confirmations map[string]Event + externalTools map[string]Event + sessionToolUses map[string]bool + toolUsesSinceStatus map[string]bool + blockingEventIDs map[string]bool + blockingEventsKnown bool + idleArmedAt time.Time + idleArmPending bool } type pendingToolEvent struct { @@ -470,9 +487,6 @@ func (s *toolRunnerState) consumeList(ctx context.Context) error { } } for { - if err := s.flushResults(ctx); err != nil { - s.runner.logger.Warn("send pending tool result failed", "err", err) - } var events []Event page := s.page listed := false @@ -505,6 +519,11 @@ func (s *toolRunnerState) consumeList(ctx context.Context) error { return err } } + if listed && !listFailed { + if err := s.flushResults(ctx); err != nil { + s.runner.logger.Warn("send pending tool result failed", "err", err) + } + } if listed && !listFailed && page == "" { // MA list events 没有可从响应事件安全推进的稳定 cursor。 // fallback 模式每轮重读完整历史，并由 seen/answered 去重， @@ -539,6 +558,7 @@ func (s *toolRunnerState) processListedEvents(ctx context.Context, events []Even if !reconcile && !seenNow { continue } + s.observeSessionState(event) if event.Type != EventTypeUserToolConfirmation { touchedIdle = true lastWasEndTurn = event.Type == EventTypeSessionStatusIdle && @@ -565,11 +585,13 @@ func (s *toolRunnerState) processListedEvents(ctx context.Context, events []Even return ErrSessionTerminated } } + s.reconcileRecoveredResults() if touchedIdle { s.disarmIdle() } for _, toolEvent := range pending { - if s.isAnswered(toolUseCallID(toolEvent.event)) { + callID := toolUseCallID(toolEvent.event) + if s.isAnswered(callID) || !s.shouldHandleToolUse(callID) { continue } if err := s.handleToolUse(ctx, toolEvent.event, toolEvent.custom); err != nil { @@ -589,6 +611,13 @@ func (s *toolRunnerState) processListedEvents(ctx context.Context, events []Even return nil } +func (s *toolRunnerState) shouldHandleToolUse(callID string) bool { + if !s.blockingEventsKnown { + return true + } + return s.blockingEventIDs[callID] || s.toolUsesSinceStatus[callID] +} + func (s *toolRunnerState) noteIdleEvent(event Event) { if event.Type == EventTypeUserToolConfirmation { return @@ -604,6 +633,8 @@ func (s *toolRunnerState) handleStreamEvent(ctx context.Context, event Event) er if !s.markEventSeen(event) { return nil } + s.observeSessionState(event) + s.reconcileRecoveredResults() s.noteIdleEvent(event) return s.handleEvent(ctx, event) } @@ -681,6 +712,7 @@ func (s *toolRunnerState) markAnswered(callID string) { s.answered[callID] = true s.processed[callID] = true delete(s.pendingResults, callID) + delete(s.recoveredResults, callID) delete(s.pendingAsk, callID) delete(s.externalTools, callID) s.maybeArmPendingIdle() @@ -716,7 +748,7 @@ func (s *toolRunnerState) releaseConfirmedToolUses(ctx context.Context) error { func (s *toolRunnerState) hasUnblockedOutstandingTool(pending []pendingToolEvent) bool { for _, toolEvent := range pending { callID := toolUseCallID(toolEvent.event) - if callID == "" || s.isAnswered(callID) { + if callID == "" || s.isAnswered(callID) || !s.shouldHandleToolUse(callID) { continue } if _, ok := s.pendingAsk[callID]; ok { @@ -736,6 +768,9 @@ func (s *toolRunnerState) handleToolUse(ctx context.Context, event Event, custom return nil } if pending := s.pendingResults[callID]; pending.ID != "" { + if s.recoveredResults[callID] { + return nil + } return s.sendResult(ctx, callID, event, custom, "", pending) } if !s.ownsTool(event, custom) { @@ -945,6 +980,9 @@ func (s *toolRunnerState) skipExternalToolUse(ctx context.Context, event Event, func (s *toolRunnerState) flushResults(ctx context.Context) error { var first error for callID, event := range s.pendingResults { + if s.recoveredResults[callID] { + continue + } req := SendEventRequest{ SessionID: s.runner.session, Event: event, @@ -967,6 +1005,70 @@ func (s *toolRunnerState) flushResults(ctx context.Context) error { return first } +func (s *toolRunnerState) observeSessionState(event Event) { + s.ensureRecoveryMaps() + switch event.Type { + case EventTypeAgentToolUse, EventTypeAgentCustomToolUse: + callID := toolUseCallID(event) + if callID != "" { + s.sessionToolUses[callID] = true + s.toolUsesSinceStatus[callID] = true + } + case EventTypeSessionStatusIdle: + s.blockingEventsKnown = true + s.blockingEventIDs = map[string]bool{} + if event.StopReasonType() == SessionStopReasonRequiresAction && event.StopReason != nil { + for _, eventID := range event.StopReason.EventIDs { + s.blockingEventIDs[eventID] = true + } + } + s.toolUsesSinceStatus = map[string]bool{} + case EventTypeSessionStatusRunning, EventTypeSessionStatusRescheduled: + s.blockingEventsKnown = true + s.blockingEventIDs = map[string]bool{} + s.toolUsesSinceStatus = map[string]bool{} + } +} + +func (s *toolRunnerState) ensureRecoveryMaps() { + if s.recoveredResults == nil { + s.recoveredResults = map[string]bool{} + } + if s.sessionToolUses == nil { + s.sessionToolUses = map[string]bool{} + } + if s.toolUsesSinceStatus == nil { + s.toolUsesSinceStatus = map[string]bool{} + } + if s.blockingEventIDs == nil { + s.blockingEventIDs = map[string]bool{} + } +} + +func (s *toolRunnerState) reconcileRecoveredResults() { + if !s.blockingEventsKnown { + return + } + for callID := range s.recoveredResults { + if s.blockingEventIDs[callID] && s.sessionToolUses[callID] { + delete(s.recoveredResults, callID) + continue + } + if s.toolUsesSinceStatus[callID] { + continue + } + delete(s.recoveredResults, callID) + delete(s.pendingResults, callID) + s.runner.logger.Warn("discard stale recovered tool result", "tool_use_id", callID) + if discarder, ok := s.runner.opts.ResultStore.(ToolResultStoreDiscarder); ok { + if err := discarder.Discard(callID); err != nil { + s.runner.logger.Warn("discard persisted tool result failed", "tool_use_id", callID, "err", err) + } + } + } + s.maybeArmPendingIdle() +} + func (s *toolRunnerState) retrySendEvent(ctx context.Context, req SendEventRequest, callID string) (bool, error) { var err error for attempt := 0; attempt < sessionRunnerSendRetries; attempt++ { diff --git a/arkruntime/selfhosted/session_tool_runner_test.go b/arkruntime/selfhosted/session_tool_runner_test.go index 3892daa..c8a5ec2 100644 --- a/arkruntime/selfhosted/session_tool_runner_test.go +++ b/arkruntime/selfhosted/session_tool_runner_test.go @@ -69,6 +69,31 @@ type runnerFailingMarkSentStore struct { markCalls int } +type runnerRecordingStore struct { + discarded []string + marked []string +} + +func (s *runnerRecordingStore) Recover() (map[string]Event, map[string]bool, error) { + return nil, nil, nil +} + +func (s *runnerRecordingStore) Begin(string, Event) (ToolCallStoreDecision, error) { + return ToolCallStoreDecision{}, nil +} + +func (s *runnerRecordingStore) SaveResult(string, Event) error { return nil } + +func (s *runnerRecordingStore) MarkSent(callID string) error { + s.marked = append(s.marked, callID) + return nil +} + +func (s *runnerRecordingStore) Discard(callID string) error { + s.discarded = append(s.discarded, callID) + return nil +} + func (s *runnerFailingMarkSentStore) Recover() (map[string]Event, map[string]bool, error) { return nil, nil, nil } @@ -287,3 +312,117 @@ func TestSessionToolRunnerConflictIsRetryable(t *testing.T) { t.Fatal("409 conflict must remain retryable") } } + +func TestSessionToolRunnerFiltersRecoveredResultsAgainstCurrentBlockers(t *testing.T) { + tool := &runnerTestTool{} + api := &runnerTestAPI{} + store := &runnerRecordingStore{} + runner := NewSessionToolRunner(context.Background(), api, "session-id", SessionToolRunnerOptions{ + CustomTools: map[string]toolset.Tool{tool.Name(): tool}, + ResultStore: store, + Logger: log.New(io.Discard, "", 0), + }) + runner.events = make(chan ToolCallResult, 1) + foreign := NewUserCustomToolResultEvent("foreign-call", nil, false, "") + stale := NewUserCustomToolResultEvent("stale-call", nil, false, "") + state := &toolRunnerState{ + runner: runner, + processed: map[string]bool{}, + seen: map[string]bool{}, + answered: map[string]bool{}, + pendingResults: map[string]Event{"foreign-call": foreign, "stale-call": stale}, + recoveredResults: map[string]bool{"foreign-call": true, "stale-call": true}, + pendingAsk: map[string]Event{}, + confirmations: map[string]Event{}, + externalTools: map[string]Event{}, + } + events := []Event{ + {ID: "stale-call", Type: EventTypeAgentCustomToolUse, Name: tool.Name(), Input: RawJSON(`{}`)}, + {ID: "current-call", Type: EventTypeAgentCustomToolUse, Name: tool.Name(), Input: RawJSON(`{}`)}, + {ID: "idle", Type: EventTypeSessionStatusIdle, StopReason: &SessionStopReason{ + Type: SessionStopReasonRequiresAction, EventIDs: []string{"current-call"}, + }}, + } + + if err := state.processListedEvents(context.Background(), events, true); err != nil { + t.Fatal(err) + } + if tool.calls != 1 || len(api.sent) != 1 || api.sent[0].CustomToolUseID != "current-call" { + t.Fatalf("tool_calls=%d sent=%+v", tool.calls, api.sent) + } + if len(state.pendingResults) != 0 || len(state.recoveredResults) != 0 { + t.Fatalf("pending=%v recovered=%v", state.pendingResults, state.recoveredResults) + } + if len(store.discarded) != 2 { + t.Fatalf("discarded=%v", store.discarded) + } +} + +func TestSessionToolRunnerResendsRecoveredCurrentBlockerWithoutExecuting(t *testing.T) { + tool := &runnerTestTool{} + api := &runnerTestAPI{} + store := &runnerRecordingStore{} + runner := NewSessionToolRunner(context.Background(), api, "session-id", SessionToolRunnerOptions{ + CustomTools: map[string]toolset.Tool{tool.Name(): tool}, + ResultStore: store, + Logger: log.New(io.Discard, "", 0), + }) + runner.events = make(chan ToolCallResult, 1) + result := NewUserCustomToolResultEvent("current-call", nil, false, "") + state := &toolRunnerState{ + runner: runner, + processed: map[string]bool{}, + seen: map[string]bool{}, + answered: map[string]bool{}, + pendingResults: map[string]Event{"current-call": result}, + recoveredResults: map[string]bool{"current-call": true}, + pendingAsk: map[string]Event{}, + confirmations: map[string]Event{}, + externalTools: map[string]Event{}, + } + + if err := state.processListedEvents(context.Background(), []Event{ + {ID: "current-call", Type: EventTypeAgentCustomToolUse, Name: tool.Name(), Input: RawJSON(`{}`)}, + {ID: "idle", Type: EventTypeSessionStatusIdle, StopReason: &SessionStopReason{ + Type: SessionStopReasonRequiresAction, EventIDs: []string{"current-call"}, + }}, + }, true); err != nil { + t.Fatal(err) + } + if tool.calls != 0 || len(api.sent) != 1 || len(store.marked) != 1 { + t.Fatalf("tool_calls=%d sent=%d marked=%v", tool.calls, len(api.sent), store.marked) + } + if len(store.discarded) != 0 { + t.Fatalf("discarded=%v", store.discarded) + } +} + +func TestSessionToolRunnerWaitsToValidateRecoveredResultUntilStatusArrives(t *testing.T) { + tool := &runnerTestTool{} + api := &runnerTestAPI{} + runner := NewSessionToolRunner(context.Background(), api, "session-id", SessionToolRunnerOptions{ + CustomTools: map[string]toolset.Tool{tool.Name(): tool}, + Logger: log.New(io.Discard, "", 0), + }) + result := NewUserCustomToolResultEvent("current-call", nil, false, "") + state := &toolRunnerState{ + runner: runner, + processed: map[string]bool{}, + seen: map[string]bool{}, + answered: map[string]bool{}, + pendingResults: map[string]Event{"current-call": result}, + recoveredResults: map[string]bool{"current-call": true}, + pendingAsk: map[string]Event{}, + confirmations: map[string]Event{}, + externalTools: map[string]Event{}, + } + + if err := state.processListedEvents(context.Background(), []Event{{ + ID: "current-call", Type: EventTypeAgentCustomToolUse, Name: tool.Name(), Input: RawJSON(`{}`), + }}, true); err != nil { + t.Fatal(err) + } + if tool.calls != 0 || len(api.sent) != 0 || !state.recoveredResults["current-call"] { + t.Fatalf("tool_calls=%d sent=%d recovered=%v", tool.calls, len(api.sent), state.recoveredResults) + } +} diff --git a/arkruntime/selfhosted/tool_result_store.go b/arkruntime/selfhosted/tool_result_store.go index f5d21bb..0b9838c 100644 --- a/arkruntime/selfhosted/tool_result_store.go +++ b/arkruntime/selfhosted/tool_result_store.go @@ -178,6 +178,22 @@ func (s *FileToolResultStore) MarkSent(callID string) error { return s.write(record) } +// Discard 删除不再属于当前 session 阻塞集合的恢复记录。 +func (s *FileToolResultStore) Discard(callID string) error { + if callID == "" { + return errors.New("call id must not be empty") + } + if err := os.Remove(s.path(callID)); err != nil && !errors.Is(err, os.ErrNotExist) { + return err + } + dir, err := os.Open(s.dir) + if err != nil { + return err + } + defer func() { _ = dir.Close() }() + return dir.Sync() +} + func (s *FileToolResultStore) read(callID string) (fileToolResultRecord, error) { return s.readPath(s.path(callID)) } @@ -255,3 +271,4 @@ func unknownToolExecutionResult(callID string, event Event) Event { } var _ ToolResultStore = (*FileToolResultStore)(nil) +var _ ToolResultStoreDiscarder = (*FileToolResultStore)(nil) diff --git a/arkruntime/selfhosted/tool_result_store_test.go b/arkruntime/selfhosted/tool_result_store_test.go index 1e97dcb..a77f16b 100644 --- a/arkruntime/selfhosted/tool_result_store_test.go +++ b/arkruntime/selfhosted/tool_result_store_test.go @@ -100,3 +100,23 @@ func TestSessionFileToolResultStoreSanitizesSessionID(t *testing.T) { t.Fatalf("store dir escaped base: %q", store.dir) } } + +func TestFileToolResultStoreDiscardRemovesRecoveredRecord(t *testing.T) { + store, err := NewSessionFileToolResultStore(t.TempDir(), "session-a") + if err != nil { + t.Fatal(err) + } + if _, err := store.Begin("call-id", Event{ID: "call-id", Type: EventTypeAgentToolUse}); err != nil { + t.Fatal(err) + } + if err := store.Discard("call-id"); err != nil { + t.Fatal(err) + } + pending, processed, err := store.Recover() + if err != nil { + t.Fatal(err) + } + if len(pending) != 0 || len(processed) != 0 { + t.Fatalf("pending=%v processed=%v", pending, processed) + } +} diff --git a/arkruntime/selfhosted/types.go b/arkruntime/selfhosted/types.go index 63fdb05..4b0cf4e 100644 --- a/arkruntime/selfhosted/types.go +++ b/arkruntime/selfhosted/types.go @@ -29,6 +29,10 @@ const ( EventTypeUserCustomToolResult = "user.custom_tool_result" // EventTypeSessionStatusIdle 表示 session 已经进入空闲状态。 EventTypeSessionStatusIdle = "session.status_idle" + // EventTypeSessionStatusRunning 表示 session 正在运行。 + EventTypeSessionStatusRunning = "session.status_running" + // EventTypeSessionStatusRescheduled 表示 session 正在重新调度。 + EventTypeSessionStatusRescheduled = "session.status_rescheduled" // EventTypeSessionStatusTerminated 表示 session 已经终止。 EventTypeSessionStatusTerminated = "session.status_terminated" // EventTypeSessionDeleted 表示 session 已经删除。