mirror of
https://gh-proxy.org/https://github.com/klinsmaya/multica.git
synced 2026-11-11 10:56:45 +08:00
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 <[email protected]> --------- Co-authored-by: Sol-Boy <[email protected]> Co-authored-by: multica-agent <[email protected]>
This commit is contained in:
importcpp
co-authored by
multica-agent
Sol-Boy
multica-agent
Sol-Boy
parent
636f8896dd
commit
9954baab86
@@ -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)
|
||||
}
|
||||
}
|
||||
}
|
||||
|
||||
|
||||
@@ -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 {
|
||||
|
||||
@@ -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)
|
||||
}
|
||||
}
|
||||
}
|
||||
@@ -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"}}'
|
||||
|
||||
Reference in New Issue
Block a user