From fe35cd4ae9c75f9ec2ea2267c9cfc3b196a450cf Mon Sep 17 00:00:00 2001 From: yushen Date: Mon, 13 Apr 2026 14:07:53 +0800 Subject: [PATCH] fix(daemon): narrow session resume retry and merge usage Address review feedback: 1. Narrow retry trigger: only retry when result.SessionID == "" (no session was established), not on any failure with PriorSessionID set 2. Merge token usage from both attempts so billing is accurate 3. Log errors when the retry itself fails to start 4. Add unit tests for mergeUsage, fallback behavior, and no-retry when session was already established Co-Authored-By: Claude Opus 4.6 (1M context) --- server/internal/daemon/daemon.go | 35 +++++- server/internal/daemon/daemon_test.go | 150 ++++++++++++++++++++++++++ 2 files changed, 182 insertions(+), 3 deletions(-) diff --git a/server/internal/daemon/daemon.go b/server/internal/daemon/daemon.go index 58608cf50..a25340698 100644 --- a/server/internal/daemon/daemon.go +++ b/server/internal/daemon/daemon.go @@ -988,12 +988,19 @@ func (d *Daemon) runTask(ctx context.Context, task Task, provider string, taskLo return TaskResult{}, err } - // Fallback: if session resume failed, retry with a fresh session. - if result.Status == "failed" && task.PriorSessionID != "" { + // Fallback: if session resume failed before establishing a session, retry + // with a fresh session. We check SessionID == "" to distinguish a resume + // failure (no session established) from a failure during actual execution. + if result.Status == "failed" && task.PriorSessionID != "" && result.SessionID == "" { + firstUsage := result.Usage taskLog.Warn("session resume failed, retrying with fresh session", "error", result.Error) execOpts.ResumeSessionID = "" - if retryResult, retryTools, retryErr := d.executeAndDrain(ctx, backend, prompt, execOpts, taskLog, task.ID); retryErr == nil { + retryResult, retryTools, retryErr := d.executeAndDrain(ctx, backend, prompt, execOpts, taskLog, task.ID) + if retryErr != nil { + taskLog.Error("fresh session also failed to start", "error", retryErr) + } else { result = retryResult + result.Usage = mergeUsage(firstUsage, result.Usage) tools = retryTools } } @@ -1182,6 +1189,28 @@ func (d *Daemon) executeAndDrain(ctx context.Context, backend agent.Backend, pro return result, toolCount.Load(), nil } +func mergeUsage(a, b map[string]agent.TokenUsage) map[string]agent.TokenUsage { + if len(a) == 0 { + return b + } + if len(b) == 0 { + return a + } + merged := make(map[string]agent.TokenUsage, len(a)+len(b)) + for model, u := range a { + merged[model] = u + } + for model, u := range b { + existing := merged[model] + existing.InputTokens += u.InputTokens + existing.OutputTokens += u.OutputTokens + existing.CacheReadTokens += u.CacheReadTokens + existing.CacheWriteTokens += u.CacheWriteTokens + merged[model] = existing + } + return merged +} + // repoDataToInfo converts daemon RepoData to repocache RepoInfo. func repoDataToInfo(repos []RepoData) []repocache.RepoInfo { info := make([]repocache.RepoInfo, len(repos)) diff --git a/server/internal/daemon/daemon_test.go b/server/internal/daemon/daemon_test.go index 868970d61..0f56b9fbe 100644 --- a/server/internal/daemon/daemon_test.go +++ b/server/internal/daemon/daemon_test.go @@ -1,9 +1,15 @@ package daemon import ( + "context" + "log/slog" "net/http" + "net/http/httptest" "strings" + "sync/atomic" "testing" + + "github.com/multica-ai/multica/server/pkg/agent" ) func TestNormalizeServerBaseURL(t *testing.T) { @@ -83,3 +89,147 @@ func TestIsWorkspaceNotFoundError(t *testing.T) { t.Fatal("did not expect 500 to be treated as workspace not found") } } + +func TestMergeUsage(t *testing.T) { + t.Parallel() + + a := map[string]agent.TokenUsage{ + "model-a": {InputTokens: 10, OutputTokens: 5}, + } + b := map[string]agent.TokenUsage{ + "model-a": {InputTokens: 20, OutputTokens: 10, CacheReadTokens: 3}, + "model-b": {InputTokens: 100}, + } + merged := mergeUsage(a, b) + + if got := merged["model-a"]; got.InputTokens != 30 || got.OutputTokens != 15 || got.CacheReadTokens != 3 { + t.Fatalf("model-a: expected {30,15,3,0}, got %+v", got) + } + if got := merged["model-b"]; got.InputTokens != 100 { + t.Fatalf("model-b: expected InputTokens=100, got %+v", got) + } + + if got := mergeUsage(nil, b); len(got) != 2 { + t.Fatal("mergeUsage(nil, b) should return b") + } + if got := mergeUsage(a, nil); len(got) != 1 { + t.Fatal("mergeUsage(a, nil) should return a") + } +} + +// fakeBackend is a test double for agent.Backend that returns preconfigured +// results. Each call to Execute pops the next entry from the results slice. +type fakeBackend struct { + calls []agent.ExecOptions + results []agent.Result + errors []error + idx atomic.Int32 +} + +func (b *fakeBackend) Execute(_ context.Context, _ string, opts agent.ExecOptions) (*agent.Session, error) { + i := int(b.idx.Add(1)) - 1 + b.calls = append(b.calls, opts) + if i < len(b.errors) && b.errors[i] != nil { + return nil, b.errors[i] + } + msgCh := make(chan agent.Message) + resCh := make(chan agent.Result, 1) + close(msgCh) + resCh <- b.results[i] + return &agent.Session{Messages: msgCh, Result: resCh}, nil +} + +func newTestDaemon(t *testing.T) *Daemon { + t.Helper() + srv := httptest.NewServer(http.HandlerFunc(func(w http.ResponseWriter, _ *http.Request) { + w.WriteHeader(http.StatusOK) + })) + t.Cleanup(srv.Close) + return &Daemon{ + client: NewClient(srv.URL), + logger: slog.Default(), + } +} + +func TestExecuteAndDrain_ResumeFailureFallback(t *testing.T) { + t.Parallel() + + d := newTestDaemon(t) + ctx := context.Background() + taskLog := slog.Default() + + fb := &fakeBackend{ + results: []agent.Result{ + {Status: "failed", Error: "session not found", Usage: map[string]agent.TokenUsage{ + "m1": {InputTokens: 5}, + }}, + {Status: "completed", Output: "done", SessionID: "new-sess", Usage: map[string]agent.TokenUsage{ + "m1": {InputTokens: 10, OutputTokens: 20}, + }}, + }, + } + + // First attempt: resume fails (no SessionID in result). + opts := agent.ExecOptions{ResumeSessionID: "stale-id"} + result, _, err := d.executeAndDrain(ctx, fb, "prompt", opts, taskLog, "task-1") + if err != nil { + t.Fatalf("first call error: %v", err) + } + if result.Status != "failed" || result.SessionID != "" { + t.Fatalf("expected failed result with empty SessionID, got %+v", result) + } + + // Simulate the retry logic from runTask. + if result.Status == "failed" && result.SessionID == "" { + firstUsage := result.Usage + opts.ResumeSessionID = "" + retryResult, _, retryErr := d.executeAndDrain(ctx, fb, "prompt", opts, taskLog, "task-1") + if retryErr != nil { + t.Fatalf("retry error: %v", retryErr) + } + result = retryResult + result.Usage = mergeUsage(firstUsage, result.Usage) + } + + if result.Status != "completed" || result.Output != "done" { + t.Fatalf("expected completed result, got %+v", result) + } + if result.SessionID != "new-sess" { + t.Fatalf("expected new-sess, got %s", result.SessionID) + } + // Usage should be merged. + if u := result.Usage["m1"]; u.InputTokens != 15 || u.OutputTokens != 20 { + t.Fatalf("expected merged usage {15,20}, got %+v", u) + } + // Second call should NOT have ResumeSessionID. + if fb.calls[1].ResumeSessionID != "" { + t.Fatal("retry should not have ResumeSessionID") + } +} + +func TestExecuteAndDrain_NoRetryWhenSessionEstablished(t *testing.T) { + t.Parallel() + + d := newTestDaemon(t) + + fb := &fakeBackend{ + results: []agent.Result{ + {Status: "failed", Error: "model error", SessionID: "valid-sess"}, + }, + } + + opts := agent.ExecOptions{ResumeSessionID: "some-id"} + result, _, err := d.executeAndDrain(context.Background(), fb, "p", opts, slog.Default(), "t") + if err != nil { + t.Fatal(err) + } + + // SessionID is set → session was established → should NOT retry. + shouldRetry := result.Status == "failed" && result.SessionID == "" + if shouldRetry { + t.Fatal("should not retry when SessionID is present") + } + if int(fb.idx.Load()) != 1 { + t.Fatalf("expected 1 call, got %d", fb.idx.Load()) + } +}