From 9954baab868ca3ef3ecbc4e1f8a0ea3e09cc9061 Mon Sep 17 00:00:00 2001 From: importcpp Date: Wed, 16 Sep 2026 15:20:06 +0800 Subject: [PATCH] MUL-7397: fix(qwen): accumulate fallback usage by message (#8443) * fix(qwen): accumulate fallback usage by message * fix(qwen): count each usage message once * test(service): scope cancelled recovery assertions to signals Co-authored-by: multica-agent --------- Co-authored-by: Sol-Boy Co-authored-by: multica-agent --- ...delegated_failure_recovery_settled_test.go | 37 ++++++- server/pkg/agent/qwen.go | 45 +++++++- server/pkg/agent/qwen_fallback_usage_test.go | 101 ++++++++++++++++++ server/pkg/agent/qwen_test.go | 5 + 4 files changed, 183 insertions(+), 5 deletions(-) create mode 100644 server/pkg/agent/qwen_fallback_usage_test.go diff --git a/server/internal/service/delegated_failure_recovery_settled_test.go b/server/internal/service/delegated_failure_recovery_settled_test.go index f3b3b7938..a9a77521e 100644 --- a/server/internal/service/delegated_failure_recovery_settled_test.go +++ b/server/internal/service/delegated_failure_recovery_settled_test.go @@ -5,6 +5,7 @@ import ( "testing" "github.com/jackc/pgx/v5/pgtype" + "github.com/multica-ai/multica/server/internal/testutil" "github.com/multica-ai/multica/server/internal/util" db "github.com/multica-ai/multica/server/pkg/db/generated" ) @@ -161,8 +162,40 @@ func TestUserCancelledRecoveryIsSettled(t *testing.T) { if !f.settled(t, recoveryCommentID) { t.Fatal("user cancellation did not settle the recovery signal") } - if result, err := svc.RecoverPendingDelegatedFailures(ctx, 100); err != nil || result != (DelegatedFailureRecoverySweepResult{}) { - t.Fatalf("sweep after user cancel = %+v, %v; want zero result, nil", result, err) + // Other packages share DATABASE_URL and can leave another workspace's + // recovery pending while this test runs the global sweep. + other, otherSvc := seedDelegatedFailureFixture(t) + otherTaskID, otherCommentID := other.seedRecoverySignal(t, otherSvc) + if _, err := otherSvc.CancelTask(ctx, otherTaskID); err != nil { + t.Fatalf("cancel other workspace recovery: %v", err) + } + if other.settled(t, otherCommentID) { + t.Fatal("undelivered system cancellation settled the other recovery") + } + if _, err := svc.RecoverPendingDelegatedFailures(ctx, 100); err != nil { + t.Fatalf("sweep after user cancel: %v", err) + } + if !f.settled(t, recoveryCommentID) { + t.Fatal("sweep reopened the user-cancelled recovery") + } + // Assert on each signal, not the sweep's totals across all workspaces. + dbfx := testutil.New(f.pool, f.workspaceID, f.userID) + for _, tc := range []struct { + name string + commentID pgtype.UUID + wantTasks int + }{ + {"user_cancelled", recoveryCommentID, 1}, + {"other_workspace_pending", otherCommentID, 2}, + } { + var taskCount int + dbfx.QueryRow(t, ` + SELECT count(*) FROM agent_task_queue + WHERE trigger_comment_id = $1 OR $1::uuid = ANY(coalesced_comment_ids) + `, tc.commentID).Scan(&taskCount) + if taskCount != tc.wantTasks { + t.Errorf("%s recovery tasks = %d, want %d", tc.name, taskCount, tc.wantTasks) + } } } diff --git a/server/pkg/agent/qwen.go b/server/pkg/agent/qwen.go index cd65eab5a..f1c03a9b7 100644 --- a/server/pkg/agent/qwen.go +++ b/server/pkg/agent/qwen.go @@ -247,6 +247,7 @@ type qwenStreamEvent struct { } type qwenMessage struct { + ID string `json:"id,omitempty"` Model string `json:"model,omitempty"` Content []qwenContentBlock `json:"content"` Usage *qwenUsage `json:"usage,omitempty"` @@ -273,6 +274,10 @@ type qwenStreamState struct { sessionID, model, lastAssistantText, finalResultText, lastEventType string sawResult, resultIsError bool usage map[string]TokenUsage + fallbackUsage map[string]TokenUsage + seenUsageMessageIDs map[string]struct{} + anonymousUsage map[string]TokenUsage + hasResultUsage bool eventCount, invalidEventCount, assistantEventCount, toolUseCount int unreadableAssistantCount int } @@ -289,7 +294,7 @@ func handleQwenEvent(event qwenStreamEvent, ch chan<- Message, state *qwenStream trySend(ch, Message{Type: MessageStatus, Status: "running", SessionID: state.sessionID}) case "assistant": state.assistantEventCount++ - turn, model := handleQwenAssistant(event.Message, ch, state.usage) + turn, model := handleQwenAssistant(event.Message, ch, state) if model != "" { state.model = model } @@ -312,6 +317,7 @@ func handleQwenEvent(event qwenStreamEvent, ch chan<- Message, state *qwenStream } if usage := qwenResultUsage(event.Usage, state.model); len(usage) > 0 { state.usage = usage + state.hasResultUsage = true } case "error": // Be fail-closed if a later Qwen release emits a terminal error event. @@ -321,7 +327,7 @@ func handleQwenEvent(event qwenStreamEvent, ch chan<- Message, state *qwenStream } } -func handleQwenAssistant(raw json.RawMessage, ch chan<- Message, usage map[string]TokenUsage) (assistantTurn, string) { +func handleQwenAssistant(raw json.RawMessage, ch chan<- Message, state *qwenStreamState) (assistantTurn, string) { var message qwenMessage if json.Unmarshal(raw, &message) != nil { // Unreadable body: understood stays false so the caller drops any @@ -330,7 +336,7 @@ func handleQwenAssistant(raw json.RawMessage, ch chan<- Message, usage map[strin } turn := assistantTurn{understood: true} if message.Usage != nil && message.Model != "" { - usage[message.Model] = qwenTokenUsage(message.Usage) + state.accumulateAssistantUsage(message) } var text strings.Builder tools := 0 @@ -365,6 +371,39 @@ func handleQwenAssistant(raw json.RawMessage, ch chan<- Message, usage map[strin return turn, message.Model } +func (s *qwenStreamState) accumulateAssistantUsage(message qwenMessage) { + if s.fallbackUsage == nil { + s.fallbackUsage = make(map[string]TokenUsage) + s.seenUsageMessageIDs = make(map[string]struct{}) + s.anonymousUsage = make(map[string]TokenUsage) + } + current := qwenTokenUsage(message.Usage) + var previous TokenUsage + if message.ID != "" { + // Qwen emits full assistant messages at finalization. Count the first + // non-empty snapshot once by ID, keeping its counters and model together. + if current.InputTokens == 0 && current.OutputTokens == 0 && current.CacheReadTokens == 0 { + return + } + if _, seen := s.seenUsageMessageIDs[message.ID]; seen { + return + } + s.seenUsageMessageIDs[message.ID] = struct{}{} + } else { + // Without an ID, preserve the existing latest-per-model best effort. + previous = s.anonymousUsage[message.Model] + s.anonymousUsage[message.Model] = current + } + total := s.fallbackUsage[message.Model] + total.InputTokens += current.InputTokens - previous.InputTokens + total.OutputTokens += current.OutputTokens - previous.OutputTokens + total.CacheReadTokens += current.CacheReadTokens - previous.CacheReadTokens + s.fallbackUsage[message.Model] = total + if !s.hasResultUsage { + s.usage = s.fallbackUsage + } +} + func handleQwenUser(raw json.RawMessage, ch chan<- Message) { var message qwenMessage if json.Unmarshal(raw, &message) != nil { diff --git a/server/pkg/agent/qwen_fallback_usage_test.go b/server/pkg/agent/qwen_fallback_usage_test.go new file mode 100644 index 000000000..90d2c8f90 --- /dev/null +++ b/server/pkg/agent/qwen_fallback_usage_test.go @@ -0,0 +1,101 @@ +package agent + +import ( + "context" + "path/filepath" + "reflect" + "testing" + "time" +) + +func qwenFallbackAssistant(t *testing.T, id, model string, usage any) qwenStreamEvent { + t.Helper() + return qwenStreamEvent{Type: "assistant", Message: mustMarshal(t, map[string]any{ + "id": id, "model": model, "usage": usage, + "content": []map[string]string{{"type": "text", "text": "visible"}}, + })} +} + +func TestQwenFallbackUsageAccumulatesMessageSnapshots(t *testing.T) { + t.Parallel() + first := qwenFallbackAssistant(t, "message-1", "test", &qwenUsage{InputTokens: 100, OutputTokens: 50}) + second := qwenFallbackAssistant(t, "message-2", "test", &qwenUsage{InputTokens: 200, OutputTokens: 20}) + final := qwenStreamEvent{Type: "result", Usage: &qwenUsage{InputTokens: 500, OutputTokens: 80}} + for _, tc := range []struct { + name string + events []qwenStreamEvent + want map[string]TokenUsage + }{ + {"two_messages", []qwenStreamEvent{first, second}, map[string]TokenUsage{"test": {InputTokens: 300, OutputTokens: 70}}}, + {"duplicate", []qwenStreamEvent{first, first}, map[string]TokenUsage{"test": {InputTokens: 100, OutputTokens: 50}}}, + {"interleaved_duplicate", []qwenStreamEvent{first, second, first}, map[string]TokenUsage{"test": {InputTokens: 300, OutputTokens: 70}}}, + {"placeholder_then_usage", []qwenStreamEvent{qwenFallbackAssistant(t, "message-1", "test", &qwenUsage{}), first}, map[string]TokenUsage{"test": {InputTokens: 100, OutputTokens: 50}}}, + {"changed_duplicate_keeps_first_snapshot", []qwenStreamEvent{first, qwenFallbackAssistant(t, "message-1", "test", &qwenUsage{InputTokens: 120, OutputTokens: 60}), first}, map[string]TokenUsage{"test": {InputTokens: 100, OutputTokens: 50}}}, + {"decreasing_counter_keeps_whole_snapshot", []qwenStreamEvent{ + qwenFallbackAssistant(t, "message-1", "test", &qwenUsage{InputTokens: 100, OutputTokens: 60, CacheReadInputTokens: 20}), + qwenFallbackAssistant(t, "message-1", "test", &qwenUsage{InputTokens: 100, OutputTokens: 50, CacheReadInputTokens: 40}), + }, map[string]TokenUsage{"test": {InputTokens: 80, OutputTokens: 60, CacheReadTokens: 20}}}, + {"corrected_model_does_not_count_message_twice", []qwenStreamEvent{first, qwenFallbackAssistant(t, "message-1", "other", &qwenUsage{InputTokens: 100, OutputTokens: 50})}, map[string]TokenUsage{"test": {InputTokens: 100, OutputTokens: 50}}}, + {"empty_snapshot_does_not_claim_model", []qwenStreamEvent{qwenFallbackAssistant(t, "message-1", "other", &qwenUsage{}), first}, map[string]TokenUsage{"test": {InputTokens: 100, OutputTokens: 50}}}, + {"missing_usage_does_not_claim_id", []qwenStreamEvent{qwenFallbackAssistant(t, "message-1", "test", nil), first}, map[string]TokenUsage{"test": {InputTokens: 100, OutputTokens: 50}}}, + {"cache_only_snapshot_claims_id", []qwenStreamEvent{qwenFallbackAssistant(t, "message-1", "test", &qwenUsage{InputTokens: 20, CacheReadInputTokens: 20}), first}, map[string]TokenUsage{"test": {CacheReadTokens: 20}}}, + {"separate_models", []qwenStreamEvent{first, qwenFallbackAssistant(t, "message-2", "other", &qwenUsage{InputTokens: 20, OutputTokens: 10})}, map[string]TokenUsage{"test": {InputTokens: 100, OutputTokens: 50}, "other": {InputTokens: 20, OutputTokens: 10}}}, + {"no_id_keeps_latest_snapshot", []qwenStreamEvent{qwenFallbackAssistant(t, "", "test", &qwenUsage{InputTokens: 100, OutputTokens: 50}), qwenFallbackAssistant(t, "", "test", &qwenUsage{InputTokens: 200, OutputTokens: 20})}, map[string]TokenUsage{"test": {InputTokens: 200, OutputTokens: 20}}}, + {"no_id_decrease_preserves_identified_usage", []qwenStreamEvent{first, qwenFallbackAssistant(t, "", "test", &qwenUsage{InputTokens: 200, OutputTokens: 20}), qwenFallbackAssistant(t, "", "test", &qwenUsage{InputTokens: 50, OutputTokens: 10})}, map[string]TokenUsage{"test": {InputTokens: 150, OutputTokens: 60}}}, + {"no_id_empty_snapshot_clears_only_anonymous_usage", []qwenStreamEvent{first, qwenFallbackAssistant(t, "", "test", &qwenUsage{InputTokens: 200, OutputTokens: 20}), qwenFallbackAssistant(t, "", "test", &qwenUsage{})}, map[string]TokenUsage{"test": {InputTokens: 100, OutputTokens: 50}}}, + {"empty_result_keeps_fallback", []qwenStreamEvent{first, second, {Type: "result", Usage: &qwenUsage{}}}, map[string]TokenUsage{"test": {InputTokens: 300, OutputTokens: 70}}}, + {"final_total_wins", []qwenStreamEvent{first, second, final}, map[string]TokenUsage{"test": {InputTokens: 500, OutputTokens: 80}}}, + {"late_assistant_cannot_replace_final_total", []qwenStreamEvent{first, final, second}, map[string]TokenUsage{"test": {InputTokens: 500, OutputTokens: 80}}}, + } { + t.Run(tc.name, func(t *testing.T) { + state := qwenStreamState{usage: make(map[string]TokenUsage)} + messages := make(chan Message, len(tc.events)) + assistantCount := 0 + for _, event := range tc.events { + handleQwenEvent(event, messages, &state) + if event.Type == "assistant" { + assistantCount++ + } + } + if !reflect.DeepEqual(state.usage, tc.want) { + t.Fatalf("usage = %+v, want %+v", state.usage, tc.want) + } + if len(messages) != assistantCount { + t.Fatalf("usage deduplication suppressed transcript messages: got %d, want %d", len(messages), assistantCount) + } + }) + } +} + +func TestQwenFallbackUsagePreservesFirstCacheSnapshot(t *testing.T) { + t.Parallel() + state := qwenStreamState{usage: make(map[string]TokenUsage)} + for _, event := range []qwenStreamEvent{ + qwenFallbackAssistant(t, "message-1", "test", &qwenUsage{InputTokens: 100, CacheReadInputTokens: 20}), + qwenFallbackAssistant(t, "message-2", "test", &qwenUsage{InputTokens: 100, CacheReadInputTokens: 30}), + qwenFallbackAssistant(t, "message-1", "test", &qwenUsage{InputTokens: 100, CacheReadInputTokens: 40}), + } { + handleQwenEvent(event, make(chan Message, 1), &state) + } + if got, want := state.usage["test"], (TokenUsage{InputTokens: 150, CacheReadTokens: 50}); got != want { + t.Fatalf("usage = %+v, want %+v", got, want) + } +} + +func TestQwenExecuteRetainsEarlierUsageOnProcessFailure(t *testing.T) { + backend := newFakeQwenBackend(t, map[string]string{"QWEN_MODE": "usage-fallback", "QWEN_STDIN_FILE": filepath.Join(t.TempDir(), "stdin")}) + // Reuse the backend to ensure snapshots belong to one Execute call. + for i := 0; i < 2; i++ { + session, err := backend.Execute(context.Background(), "test", ExecOptions{Timeout: 5 * time.Second}) + if err != nil { + t.Fatal(err) + } + _, result := awaitQwenResult(t, session) + if result.Status != "failed" { + t.Fatalf("status = %s, want failed", result.Status) + } + if got := result.Usage["qwen-test"]; got != (TokenUsage{InputTokens: 300, OutputTokens: 70}) { + t.Fatalf("usage = %+v, want 300 input / 70 output", got) + } + } +} diff --git a/server/pkg/agent/qwen_test.go b/server/pkg/agent/qwen_test.go index cce8f5c23..f28b067d3 100644 --- a/server/pkg/agent/qwen_test.go +++ b/server/pkg/agent/qwen_test.go @@ -85,6 +85,11 @@ while [ "$#" -gt 0 ]; do shift done case "$QWEN_MODE" in + usage-fallback) + printf '%s\n' '{"type":"assistant","message":{"id":"message-1","model":"qwen-test","content":[{"type":"text","text":"first"}],"usage":{"input_tokens":100,"output_tokens":50}}}' + printf '%s\n' '{"type":"assistant","message":{"id":"message-2","model":"qwen-test","content":[{"type":"text","text":"second"}],"usage":{"input_tokens":200,"output_tokens":20}}}' + exit 7 + ;; error) printf '%s\n' '{"type":"system","subtype":"init","session_id":"sess-error","model":"qwen-test"}' printf '%s\n' '{"type":"result","subtype":"error_during_execution","session_id":"sess-error","is_error":true,"error":{"type":"authentication_error","message":"synthetic Qwen authentication failure"}}'