diff --git a/agent/claudecode/session.go b/agent/claudecode/session.go index a41637a..33f6cc4 100644 --- a/agent/claudecode/session.go +++ b/agent/claudecode/session.go @@ -292,7 +292,25 @@ func (cs *claudeSession) handleResult(raw map[string]any) { if sid, ok := raw["session_id"].(string); ok && sid != "" { cs.sessionID.Store(sid) } - evt := core.Event{Type: core.EventResult, Content: content, SessionID: cs.CurrentSessionID(), Done: true} + + var inputTokens, outputTokens int + if usage, ok := raw["usage"].(map[string]any); ok { + if v, ok := usage["input_tokens"].(float64); ok { + inputTokens = int(v) + } + if v, ok := usage["output_tokens"].(float64); ok { + outputTokens = int(v) + } + } + + evt := core.Event{ + Type: core.EventResult, + Content: content, + SessionID: cs.CurrentSessionID(), + Done: true, + InputTokens: inputTokens, + OutputTokens: outputTokens, + } select { case cs.events <- evt: case <-cs.ctx.Done(): diff --git a/agent/claudecode/session_test.go b/agent/claudecode/session_test.go new file mode 100644 index 0000000..3fa8103 --- /dev/null +++ b/agent/claudecode/session_test.go @@ -0,0 +1,67 @@ +package claudecode + +import ( + "context" + "testing" + + "github.com/chenhg5/cc-connect/core" +) + +func TestHandleResultParsesUsage(t *testing.T) { + ctx, cancel := context.WithCancel(context.Background()) + defer cancel() + + cs := &claudeSession{ + events: make(chan core.Event, 8), + ctx: ctx, + } + cs.sessionID.Store("test-session") + cs.alive.Store(true) + + raw := map[string]any{ + "type": "result", + "result": "done", + "session_id": "test-session", + "usage": map[string]any{ + "input_tokens": float64(150000), + "output_tokens": float64(2000), + }, + } + + cs.handleResult(raw) + + evt := <-cs.events + if evt.InputTokens != 150000 { + t.Errorf("InputTokens = %d, want 150000", evt.InputTokens) + } + if evt.OutputTokens != 2000 { + t.Errorf("OutputTokens = %d, want 2000", evt.OutputTokens) + } +} + +func TestHandleResultNoUsage(t *testing.T) { + ctx, cancel := context.WithCancel(context.Background()) + defer cancel() + + cs := &claudeSession{ + events: make(chan core.Event, 8), + ctx: ctx, + } + cs.sessionID.Store("test-session") + cs.alive.Store(true) + + raw := map[string]any{ + "type": "result", + "result": "done", + } + + cs.handleResult(raw) + + evt := <-cs.events + if evt.InputTokens != 0 { + t.Errorf("InputTokens = %d, want 0", evt.InputTokens) + } + if evt.OutputTokens != 0 { + t.Errorf("OutputTokens = %d, want 0", evt.OutputTokens) + } +} diff --git a/core/engine.go b/core/engine.go index 9681445..10cb20d 100644 --- a/core/engine.go +++ b/core/engine.go @@ -14,6 +14,7 @@ import ( "sort" "strconv" "strings" + "regexp" "sync" "sync/atomic" "time" @@ -1669,8 +1670,16 @@ func (e *Engine) processInteractiveMessageWith(p Platform, msg *Message, session if agent != e.agent { agentOverride = agent } + // Per-workspace connectedOnce flag so that idle-reaped workspaces + // re-use --continue on their first reconnection. + var wsConnectedOnce *atomic.Bool + if workspaceDir != "" { + if ws := e.workspacePool.Get(workspaceDir); ws != nil { + wsConnectedOnce = &ws.hasConnectedOnce + } + } consumeBridge := messageConsumesFirstContinueBridge(msg) - state := e.getOrCreateInteractiveStateWith(interactiveKey, p, msg.ReplyCtx, session, sessions, agentOverride, ccSessionKey, consumeBridge) + state := e.getOrCreateInteractiveStateWith(interactiveKey, p, msg.ReplyCtx, session, sessions, agentOverride, ccSessionKey, consumeBridge, wsConnectedOnce) // Set workspaceDir on the state for idle reaper identification if workspaceDir != "" { @@ -1731,7 +1740,7 @@ func (e *Engine) processInteractiveMessageWith(p Platform, msg *Message, session e.cleanupInteractiveState(interactiveKey, state) e.send(p, msg.ReplyCtx, e.i18n.T(MsgSessionRestarting)) - state = e.getOrCreateInteractiveStateWith(interactiveKey, p, msg.ReplyCtx, session, sessions, agentOverride, ccSessionKey, consumeBridge) + state = e.getOrCreateInteractiveStateWith(interactiveKey, p, msg.ReplyCtx, session, sessions, agentOverride, ccSessionKey, consumeBridge, wsConnectedOnce) if workspaceDir != "" { state.mu.Lock() state.workspaceDir = workspaceDir @@ -1844,13 +1853,14 @@ func messageConsumesFirstContinueBridge(msg *Message) bool { } } -// getOrCreateInteractiveStateWith accepts -// an optional agent override for multi-workspace mode. When agentOverride is non-nil -// it is used instead of e.agent to start the session. +// getOrCreateInteractiveStateWith accepts an optional agent override for multi-workspace mode. +// When agentOverride is non-nil it is used instead of e.agent to start the session. // ccSessionKey, when non-empty, is used for CC_SESSION_KEY env injection; otherwise sessionKey is used. // consumeFirstUserContinueBridge controls whether this start may use the global // one-time ContinueSession bridge (see messageConsumesFirstContinueBridge). -func (e *Engine) getOrCreateInteractiveStateWith(sessionKey string, p Platform, replyCtx any, session *Session, sessions *SessionManager, agentOverride Agent, ccSessionKey string, consumeFirstUserContinueBridge bool) *interactiveState { +// connectedOnce, when non-nil, is used instead of the engine-level hasConnectedOnce flag +// (for per-workspace tracking after idle reap). +func (e *Engine) getOrCreateInteractiveStateWith(sessionKey string, p Platform, replyCtx any, session *Session, sessions *SessionManager, agentOverride Agent, ccSessionKey string, consumeFirstUserContinueBridge bool, connectedOnce *atomic.Bool) *interactiveState { e.interactiveMu.Lock() defer e.interactiveMu.Unlock() @@ -1932,11 +1942,15 @@ func (e *Engine) getOrCreateInteractiveStateWith(sessionKey string, p Platform, return state } - // On first real-user connection after engine startup, use --continue to pick up - // the most recent CLI session (bridges direct CLI and cc-connect usage). + // On first real-user connection after engine startup (or workspace re-creation + // after idle reap), use --continue to pick up the most recent CLI session. // Subsequent connections respect the stored session ID (or "" for fresh). startSessionID := session.GetAgentSessionID() - if consumeFirstUserContinueBridge && !e.hasConnectedOnce.Swap(true) { + once := &e.hasConnectedOnce + if connectedOnce != nil { + once = connectedOnce + } + if consumeFirstUserContinueBridge && !once.Swap(true) { startSessionID = ContinueSession } @@ -1945,10 +1959,25 @@ func (e *Engine) getOrCreateInteractiveStateWith(sessionKey string, p Platform, agentSession, err := agent.StartSession(e.ctx, startSessionID) startElapsed := time.Since(startAt) if err != nil { - slog.Error("failed to start interactive session", "error", err, "elapsed", startElapsed) - state = &interactiveState{platform: p, replyCtx: replyCtx, quiet: quietMode} - e.interactiveStates[sessionKey] = state - return state + // If resume/continue failed, try a fresh session as fallback. + if startSessionID != "" { + slog.Error("session resume failed, falling back to fresh session", + "session_key", sessionKey, "failed_session_id", startSessionID, + "error", err, "elapsed", startElapsed) + startAt = time.Now() + agentSession, err = agent.StartSession(e.ctx, "") + startElapsed = time.Since(startAt) + if err == nil { + slog.Info("fresh session started after resume failure", + "session_key", sessionKey, "elapsed", startElapsed) + } + } + if err != nil { + slog.Error("failed to start interactive session", "error", err, "elapsed", startElapsed) + state = &interactiveState{platform: p, replyCtx: replyCtx, quiet: quietMode} + e.interactiveStates[sessionKey] = state + return state + } } if startElapsed >= slowAgentStart { slog.Warn("slow agent session start", "elapsed", startElapsed, "agent", agent.Name(), "session_id", startSessionID) @@ -2281,10 +2310,16 @@ func (e *Engine) processInteractiveEvents(state *interactiveState, session *Sess fullResponse = e.i18n.T(MsgEmptyResponse) } + // Context usage indicator: prefer SDK tokens, fall back to self-reported. + sdkPlausible := event.InputTokens >= 100 + selfPct := parseSelfReportedCtx(fullResponse) + cleanResponse := ctxSelfReportRe.ReplaceAllString(fullResponse, "") + cleanResponse = strings.TrimRight(cleanResponse, "\n ") + // Evaluate auto-compress trigger (token estimate on user+assistant text, // including this turn's assistant reply before it is appended to history). if e.autoCompressEnabled && e.autoCompressMaxTokens > 0 { - estimate := estimateTokensWithPendingAssistant(session.GetHistory(0), fullResponse) + estimate := estimateTokensWithPendingAssistant(session.GetHistory(0), cleanResponse) now := time.Now() state.mu.Lock() last := state.lastAutoCompressAt @@ -2297,9 +2332,16 @@ func (e *Engine) processInteractiveEvents(state *interactiveState, session *Sess } } - session.AddHistory("assistant", fullResponse) + session.AddHistory("assistant", cleanResponse) sessions.Save() + if sdkPlausible { + cleanResponse += contextIndicator(event.InputTokens) + } else if selfPct > 0 { + cleanResponse += fmt.Sprintf("\n[ctx: ~%d%%]", selfPct) + } + fullResponse = cleanResponse + turnDuration := time.Since(turnStart) slog.Info("turn complete", "session", session.ID, @@ -2308,6 +2350,8 @@ func (e *Engine) processInteractiveEvents(state *interactiveState, session *Sess "tools", toolCount, "response_len", len(fullResponse), "turn_duration", turnDuration, + "input_tokens", event.InputTokens, + "output_tokens", event.OutputTokens, ) replyStart := time.Now() @@ -8932,3 +8976,37 @@ func gitClone(repoURL, dest string) error { } return nil } + +// ── Context usage indicator ────────────────────────────────── + +const modelContextWindow = 200_000 // Claude's context window size in tokens + +// contextIndicator returns a suffix like "\n[ctx: ~42%]" based on SDK-reported input tokens. +func contextIndicator(inputTokens int) string { + if inputTokens <= 0 { + return "" + } + pct := inputTokens * 100 / modelContextWindow + if pct > 100 { + pct = 100 + } + return fmt.Sprintf("\n[ctx: ~%d%%]", pct) +} + +// ctxSelfReportRe matches agent self-reported context lines like "[ctx: ~42%]". +var ctxSelfReportRe = regexp.MustCompile(`(?m)\n?\[ctx: ~\d+%\]`) + +// parseSelfReportedCtx extracts the percentage from a self-reported "[ctx: ~XX%]" line. +func parseSelfReportedCtx(s string) int { + m := ctxSelfReportRe.FindString(s) + if m == "" { + return 0 + } + start := strings.Index(m, "~") + 1 + end := strings.Index(m, "%") + if start <= 0 || end <= start { + return 0 + } + v, _ := strconv.Atoi(m[start:end]) + return v +} diff --git a/core/engine_test.go b/core/engine_test.go index b8a4513..665e4ce 100644 --- a/core/engine_test.go +++ b/core/engine_test.go @@ -3505,7 +3505,7 @@ func TestFirstContinueBridge_SyntheticDoesNotConsume(t *testing.T) { t.Fatal("TryLock") } - e.getOrCreateInteractiveStateWith(key, p, "ctx", sess, e.sessions, nil, "", false) + e.getOrCreateInteractiveStateWith(key, p, "ctx", sess, e.sessions, nil, "", false, nil) if got := ag.startArgs(); len(got) != 1 || got[0] != "" { t.Fatalf("synthetic start ids = %#v want [\"\"]", got) } @@ -3515,7 +3515,7 @@ func TestFirstContinueBridge_SyntheticDoesNotConsume(t *testing.T) { e.cleanupInteractiveState(key) - e.getOrCreateInteractiveStateWith(key, p, "ctx", sess, e.sessions, nil, "", true) + e.getOrCreateInteractiveStateWith(key, p, "ctx", sess, e.sessions, nil, "", true, nil) got := ag.startArgs() if len(got) != 2 { t.Fatalf("want 2 StartSession calls, got %#v", got) @@ -3625,7 +3625,7 @@ func TestSessionMismatch_RecyclesStaleAgent(t *testing.T) { // The active Session now wants a DIFFERENT agent session ID. session := &Session{AgentSessionID: "new-agent-id"} - state := e.getOrCreateInteractiveStateWith(key, p, "ctx", session, e.sessions, nil, "", true) + state := e.getOrCreateInteractiveStateWith(key, p, "ctx", session, e.sessions, nil, "", true, nil) if state.agentSession == oldSess { t.Fatal("expected stale agent session to be replaced") @@ -3664,7 +3664,7 @@ func TestSessionMismatch_DoesNotLeakQuiet(t *testing.T) { // Active session wants "new-id", which mismatches "old-id". session := &Session{AgentSessionID: "new-id"} - state := e.getOrCreateInteractiveStateWith(key, p, "ctx", session, e.sessions, nil, "", true) + state := e.getOrCreateInteractiveStateWith(key, p, "ctx", session, e.sessions, nil, "", true, nil) state.mu.Lock() q := state.quiet @@ -3695,7 +3695,7 @@ func TestSessionMismatch_ReusesWhenIDsMatch(t *testing.T) { session := &Session{AgentSessionID: "matching-id"} - state := e.getOrCreateInteractiveStateWith(key, p, "ctx", session, e.sessions, nil, "", true) + state := e.getOrCreateInteractiveStateWith(key, p, "ctx", session, e.sessions, nil, "", true, nil) if state != existingState { t.Fatal("expected existing state to be reused when session IDs match") } @@ -3713,7 +3713,7 @@ func TestSessionIDWriteback_ImmediateAfterStartSession(t *testing.T) { key := "test:user1" session := &Session{AgentSessionID: ""} // empty — no prior binding - e.getOrCreateInteractiveStateWith(key, p, "ctx", session, e.sessions, nil, "", true) + e.getOrCreateInteractiveStateWith(key, p, "ctx", session, e.sessions, nil, "", true, nil) got := session.GetAgentSessionID() @@ -3733,7 +3733,7 @@ func TestSessionIDWriteback_DoesNotOverwriteExisting(t *testing.T) { key := "test:user1" session := &Session{AgentSessionID: "existing-uuid"} - e.getOrCreateInteractiveStateWith(key, p, "ctx", session, e.sessions, nil, "", true) + e.getOrCreateInteractiveStateWith(key, p, "ctx", session, e.sessions, nil, "", true, nil) got := session.GetAgentSessionID() @@ -3769,7 +3769,7 @@ func TestStaleGoroutineCleanup_RaceSimulation(t *testing.T) { // Step 3: New turn creates Session B and calls getOrCreateInteractiveStateWith. sessionB := &Session{AgentSessionID: ""} - newState := e.getOrCreateInteractiveStateWith(key, p, "ctx", sessionB, e.sessions, nil, "", true) + newState := e.getOrCreateInteractiveStateWith(key, p, "ctx", sessionB, e.sessions, nil, "", true, nil) // Verify S2 is in the map. e.interactiveMu.Lock() @@ -4047,6 +4047,136 @@ func TestCmdBindSetup_UsesSharedLogic(t *testing.T) { } } +// --- session resilience tests --- + +// stubStartSessionAgent records StartSession calls and can fail on specific session IDs. +type stubStartSessionAgent struct { + calls []string + failIDs map[string]error // session IDs that should fail + mu sync.Mutex +} + +func (a *stubStartSessionAgent) Name() string { return "stub" } +func (a *stubStartSessionAgent) StartSession(_ context.Context, sessionID string) (AgentSession, error) { + a.mu.Lock() + a.calls = append(a.calls, sessionID) + a.mu.Unlock() + + if err, ok := a.failIDs[sessionID]; ok { + return nil, err + } + return &stubAgentSession{}, nil +} +func (a *stubStartSessionAgent) ListSessions(_ context.Context) ([]AgentSessionInfo, error) { + return nil, nil +} +func (a *stubStartSessionAgent) Stop() error { return nil } + +func TestResumeFailureFallbackToFreshSession(t *testing.T) { + agent := &stubStartSessionAgent{ + failIDs: map[string]error{ + ContinueSession: fmt.Errorf("Prompt is too long"), + }, + } + + ctx, cancel := context.WithCancel(context.Background()) + defer cancel() + + e := &Engine{ + agent: agent, + sessions: NewSessionManager(""), + ctx: ctx, + i18n: NewI18n("en"), + interactiveStates: make(map[string]*interactiveState), + display: DisplayCfg{}, + } + + session := e.sessions.GetOrCreateActive("test:user1") + session.SetAgentSessionID("old-session-id", "stub") + + p := &stubPlatformEngine{n: "test"} + state := e.getOrCreateInteractiveStateWith("test:user1", p, "ctx", session, e.sessions, nil, "", true, nil) + + if state.agentSession == nil { + t.Fatal("expected agentSession to be non-nil after fallback") + } + + agent.mu.Lock() + calls := append([]string{}, agent.calls...) + agent.mu.Unlock() + + if len(calls) != 2 { + t.Fatalf("expected 2 StartSession calls, got %d: %v", len(calls), calls) + } + // First call should be ContinueSession (first connection uses --continue) + if calls[0] != ContinueSession { + t.Fatalf("first StartSession call = %q, want %q", calls[0], ContinueSession) + } + // Second call should be empty string (fresh session fallback) + if calls[1] != "" { + t.Fatalf("second StartSession call = %q, want empty string", calls[1]) + } +} + +func TestFreshSessionRespectedAfterFirstConnection(t *testing.T) { + agent := &stubStartSessionAgent{} + + ctx, cancel := context.WithCancel(context.Background()) + defer cancel() + + e := &Engine{ + agent: agent, + sessions: NewSessionManager(""), + ctx: ctx, + i18n: NewI18n("en"), + interactiveStates: make(map[string]*interactiveState), + display: DisplayCfg{}, + } + // Simulate first connection already happened + e.hasConnectedOnce.Store(true) + + // Create a session with no saved agent session ID (fresh session via /new) + session := e.sessions.GetOrCreateActive("test:user2") + + p := &stubPlatformEngine{n: "test"} + state := e.getOrCreateInteractiveStateWith("test:user2", p, "ctx", session, e.sessions, nil, "", true, nil) + + if state.agentSession == nil { + t.Fatal("expected agentSession to be non-nil") + } + + agent.mu.Lock() + calls := append([]string{}, agent.calls...) + agent.mu.Unlock() + + if len(calls) != 1 { + t.Fatalf("expected 1 StartSession call, got %d: %v", len(calls), calls) + } + // Should be empty string (fresh session), NOT ContinueSession + if calls[0] != "" { + t.Fatalf("StartSession call = %q, want empty string (fresh session)", calls[0]) + } +} + +func TestParseSelfReportedCtx(t *testing.T) { + tests := []struct { + input string + want int + }{ + {"here is my response\n[ctx: ~42%]", 42}, + {"no context here", 0}, + {"response\n[ctx: ~100%]", 100}, + {"response\n[ctx: ~5%]", 5}, + {"", 0}, + } + for _, tt := range tests { + got := parseSelfReportedCtx(tt.input) + if got != tt.want { + t.Errorf("parseSelfReportedCtx(%q) = %d, want %d", tt.input, got, tt.want) + } + } +} + func TestDrainEventsClosedChannel(t *testing.T) { ch := make(chan Event, 2) ch <- Event{Type: EventToolUse, Content: "a"} diff --git a/core/message.go b/core/message.go index a741a49..a2727d2 100644 --- a/core/message.go +++ b/core/message.go @@ -155,11 +155,6 @@ const ( EventThinking EventType = "thinking" // thinking/processing status ) -// ContinueSession is a sentinel value passed to Agent.StartSession to indicate -// that the agent should pick up the most recent session in the workspace -// (equivalent to `claude --continue`), rather than resuming a specific session ID. -const ContinueSession = "__continue__" - // UserQuestion represents a structured question from AskUserQuestion. type UserQuestion struct { Question string `json:"question"` @@ -187,6 +182,8 @@ type Event struct { Questions []UserQuestion // populated when ToolName == "AskUserQuestion" Done bool Error error + InputTokens int // token usage from agent result events + OutputTokens int } // HistoryEntry is one turn in a conversation. diff --git a/core/session.go b/core/session.go index a79c448..a04b14c 100644 --- a/core/session.go +++ b/core/session.go @@ -10,6 +10,10 @@ import ( "time" ) +// ContinueSession is a sentinel value for AgentSessionID that tells the agent +// to use --continue (resume most recent session) instead of a specific session ID. +const ContinueSession = "__continue__" + // Session tracks one conversation between a user and the agent. type Session struct { ID string `json:"id"` diff --git a/core/workspace_state.go b/core/workspace_state.go index 2ce0371..9e8d19d 100644 --- a/core/workspace_state.go +++ b/core/workspace_state.go @@ -4,6 +4,7 @@ import ( "log/slog" "path/filepath" "sync" + "sync/atomic" "time" ) @@ -25,11 +26,12 @@ func normalizeWorkspacePath(path string) string { // workspaceState holds the runtime state for a single workspace. type workspaceState struct { - mu sync.Mutex - workspace string - sessions *SessionManager - agent Agent - lastActivity time.Time + mu sync.Mutex + workspace string + sessions *SessionManager + agent Agent + lastActivity time.Time + hasConnectedOnce atomic.Bool // first connection after (re)creation uses --continue } func newWorkspaceState(workspace string) *workspaceState {