Skip to content

fix(selfHosted): filter recovered tool results by session blockers #26

New issue

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

Sign up for GitHub

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

Already on GitHub? Sign in to your account

Merged
ark-hand merged 1 commit into from
Sep 17, 2026
Merged

fix(selfHosted): filter recovered tool results by session blockers #26

Show file tree
Hide file tree
Changes from all commits
Commits
Show all changes
1 commit
Select commit
File filter

Filter by extension

Filter by extension
Conversations
Failed to load comments.
Loading
Jump to
Jump to file
Failed to load files.
Loading
Diff view
Diff view
30 changes: 20 additions & 10 deletions arkruntime/lib/environments/integration_test.go
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -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
}

Expand All @@ -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,
Expand Down Expand Up @@ -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"},
},
},
},
}
Expand Down Expand Up @@ -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"},
},
},
},
}
Expand Down
152 changes: 127 additions & 25 deletions arkruntime/selfhosted/session_tool_runner.go
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -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
Expand Down Expand Up @@ -180,22 +185,29 @@ 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()
if err != nil {
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
Expand All @@ -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 {
Expand Down Expand Up @@ -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
Expand Down Expand Up @@ -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 去重，
Expand Down Expand Up @@ -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 &&
Expand All @@ -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 {
Expand All @@ -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
Expand All @@ -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)
}
Expand Down Expand Up @@ -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()
Expand Down Expand Up @@ -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 {
Expand All @@ -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) {
Expand Down Expand Up @@ -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,
Expand All @@ -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++ {
Expand Down
Loading
Loading