diff --git a/pkg/agent/loop.go b/pkg/agent/loop.go index 81b979490..e480b352f 100644 --- a/pkg/agent/loop.go +++ b/pkg/agent/loop.go @@ -1451,8 +1451,10 @@ func (al *AgentLoop) processSystemMessage( return "", fmt.Errorf("no default agent for system message") } - // Use the origin session for context - sessionKey := routing.BuildAgentMainSessionKey(agent.ID) + sessionKey := msg.SessionKey + if sessionKey == "" { + sessionKey = routing.BuildAgentMainSessionKey(agent.ID) + } return al.runAgentLoop(ctx, agent, processOptions{ SessionKey: sessionKey, @@ -2369,11 +2371,12 @@ turnLoop: pubCtx, pubCancel := context.WithTimeout(context.Background(), 5*time.Second) defer pubCancel() - _ = al.bus.PublishInbound(pubCtx, bus.InboundMessage{ - Channel: "system", - SenderID: fmt.Sprintf("async:%s", asyncToolName), - ChatID: fmt.Sprintf("%s:%s", ts.channel, ts.chatID), - Content: content, + al.bus.PublishInbound(pubCtx, bus.InboundMessage{ + Channel: "system", + SenderID: fmt.Sprintf("async:%s", asyncToolName), + ChatID: fmt.Sprintf("%s:%s", ts.channel, ts.chatID), + Content: content, + SessionKey: ts.sessionKey, }) } diff --git a/pkg/agent/loop_test.go b/pkg/agent/loop_test.go index 2366b1277..4eca40499 100644 --- a/pkg/agent/loop_test.go +++ b/pkg/agent/loop_test.go @@ -2926,3 +2926,49 @@ func TestProcessMessage_ContextOverflow_AnthropicStyle(t *testing.T) { t.Fatalf("expected 2 calls for retry, got %d", provider.calls) } } + +// TestProcessSystemMessage_UsesOriginSessionKey verifies that when an async tool +// completes and publishes a system message with a SessionKey, the resulting turn +// runs in that session rather than the default "agent:main:main" session. +func TestProcessSystemMessage_UsesOriginSessionKey(t *testing.T) { + al := NewAgentLoop(&config.Config{ + Agents: config.AgentsConfig{ + Defaults: config.AgentDefaults{ + Workspace: t.TempDir(), + ModelName: "test-model", + MaxTokens: 4096, + MaxToolIterations: 10, + }, + }, + }, bus.NewMessageBus(), &mockProvider{}) + + const wantSessionKey = "agent:main:weixin:direct:testchatid" + + sub := al.SubscribeEvents(32) + defer al.UnsubscribeEvents(sub.ID) + + done := make(chan error, 1) + go func() { + _, err := al.processMessage(context.Background(), bus.InboundMessage{ + Channel: "system", + SenderID: "async:spawn", + ChatID: "weixin:testchatid", + Content: "task completed", + SessionKey: wantSessionKey, + }) + done <- err + }() + + evt := waitForEvent(t, sub.C, 5*time.Second, func(e Event) bool { + return e.Kind == EventKindTurnStart + }) + + if evt.Meta.SessionKey != wantSessionKey { + t.Errorf("turn started with session_key=%q, want %q (got agent:main:main means bug is present)", + evt.Meta.SessionKey, wantSessionKey) + } + + if err := <-done; err != nil { + t.Logf("processMessage returned error (acceptable in test env): %v", err) + } +}