diff --git a/apps/backend/internal/agent/runtime/lifecycle/event_types.go b/apps/backend/internal/agent/runtime/lifecycle/event_types.go index 8d884b0511..1a12392b56 100644 --- a/apps/backend/internal/agent/runtime/lifecycle/event_types.go +++ b/apps/backend/internal/agent/runtime/lifecycle/event_types.go @@ -129,6 +129,7 @@ type AgentStreamEventData struct { ProviderError *streams.ProviderError `json:"provider_error,omitempty"` SessionStatus string `json:"session_status,omitempty"` // "resumed" or "new" for session_status events PromptGeneration uint64 `json:"prompt_generation,omitempty"` + TurnID string `json:"turn_id,omitempty"` Data interface{} `json:"data,omitempty"` // ParentToolCallID identifies the parent Task tool call when this event @@ -662,14 +663,23 @@ func (p SessionTodosEventPayload) GetSessionID() string { // AgentID is the lifecycle execution.ID (UUID); AgentType is the CLI engine // slug (claude-acp, codex-acp, ...). The office cost subscriber derives the // provider name from AgentType — AgentID is kept for legacy consumers. +// +// TurnID is captured on the lifecycle completion frame before AgentReady can +// admit a successor prompt and is forwarded by the orchestrator. The +// orchestrator mints UsageEventID once at publish time — rather than per +// consumer — so it remains a stable idempotency key across event redelivery. +// Both are empty when unavailable (e.g. no active turn), never a synthesized +// placeholder. type SessionPromptUsageEventPayload struct { - TaskID string `json:"task_id"` - SessionID string `json:"session_id"` - AgentID string `json:"agent_id"` - AgentType string `json:"agent_type,omitempty"` - Model string `json:"model,omitempty"` - Usage *streams.PromptUsage `json:"usage"` - Timestamp string `json:"timestamp"` + TaskID string `json:"task_id"` + SessionID string `json:"session_id"` + AgentID string `json:"agent_id"` + AgentType string `json:"agent_type,omitempty"` + Model string `json:"model,omitempty"` + Usage *streams.PromptUsage `json:"usage"` + Timestamp string `json:"timestamp"` + TurnID string `json:"turn_id,omitempty"` + UsageEventID string `json:"usage_event_id,omitempty"` } // GetSessionID returns the session ID for this event (used by event routing). diff --git a/apps/backend/internal/agent/runtime/lifecycle/events.go b/apps/backend/internal/agent/runtime/lifecycle/events.go index db7b9f67b6..ab4a652a69 100644 --- a/apps/backend/internal/agent/runtime/lifecycle/events.go +++ b/apps/backend/internal/agent/runtime/lifecycle/events.go @@ -187,6 +187,7 @@ func (p *EventPublisher) PublishAgentStreamEvent(execution *AgentExecution, even ProviderError: event.ProviderError, SessionStatus: event.SessionStatus, PromptGeneration: event.PromptGeneration, + TurnID: event.TurnID, Data: event.Data, Normalized: event.NormalizedPayload, AvailableCommands: event.AvailableCommands, diff --git a/apps/backend/internal/agent/runtime/lifecycle/executor_backend.go b/apps/backend/internal/agent/runtime/lifecycle/executor_backend.go index 4818cc34e5..8e9d95fe29 100644 --- a/apps/backend/internal/agent/runtime/lifecycle/executor_backend.go +++ b/apps/backend/internal/agent/runtime/lifecycle/executor_backend.go @@ -344,6 +344,7 @@ type ExecutorCreateRequest struct { TaskEnvironmentID string // Env this execution belongs to (shared across sessions in same task) AgentProfileID string OfficeAgentProfileID string + PromptTurnID string WorkspacePath string WorkspaceSourceRoots []string Protocol string @@ -446,6 +447,7 @@ func (ri *ExecutorInstance) ToAgentExecution(req *ExecutorCreateRequest) *AgentE TaskEnvironmentID: req.TaskEnvironmentID, AgentProfileID: req.AgentProfileID, OfficeAgentProfileID: req.OfficeAgentProfileID, + promptTurnID: req.PromptTurnID, AgentID: agentID, ContainerID: ri.ContainerID, ContainerIP: ri.ContainerIP, diff --git a/apps/backend/internal/agent/runtime/lifecycle/manager_events.go b/apps/backend/internal/agent/runtime/lifecycle/manager_events.go index 589ca2e76c..9eb23e2579 100644 --- a/apps/backend/internal/agent/runtime/lifecycle/manager_events.go +++ b/apps/backend/internal/agent/runtime/lifecycle/manager_events.go @@ -299,6 +299,11 @@ func setProviderError(execution *AgentExecution, providerError *streams.Provider // handleCompleteEvent handles a "complete" agent event: flushes buffers, marks state, and signals SendPrompt. func (m *Manager) handleCompleteEvent(execution *AgentExecution, event *agentctl.AgentEvent) bool { + if event.TurnID == "" { + // Snapshot before publishing AgentReady. A queued successor may bind a + // new turn while the complete stream frame is still crossing the bus. + event.TurnID = execution.promptTurnIDSnapshot() + } isError, stopReason := completeEventResult(event) claim, claimed := m.claimPromptCompletion(execution, event, isError) if !claimed { diff --git a/apps/backend/internal/agent/runtime/lifecycle/manager_events_test.go b/apps/backend/internal/agent/runtime/lifecycle/manager_events_test.go index ee166e025b..42aa2da564 100644 --- a/apps/backend/internal/agent/runtime/lifecycle/manager_events_test.go +++ b/apps/backend/internal/agent/runtime/lifecycle/manager_events_test.go @@ -120,6 +120,35 @@ func TestHandleAgentEvent_UserMessageChunkNotBufferedAsAssistant(t *testing.T) { } } +func TestHandleAgentEvent_CompleteCarriesPromptTurnID(t *testing.T) { + mgr, eventBus := createTestManagerWithTracking() + execution := createTestExecution("exec-turn-id", "task-1", "session-1") + execution.promptTurnID = "turn-1" + if err := mgr.executionStore.Add(execution); err != nil { + t.Fatalf("add execution: %v", err) + } + generation, err := mgr.executionStore.BeginPrompt(execution.ID) + if err != nil { + t.Fatalf("begin prompt: %v", err) + } + + mgr.handleAgentEvent(execution, agentctl.AgentEvent{ + Type: streams.EventTypeComplete, + SessionID: execution.SessionID, + PromptGeneration: generation, + }) + + for _, payload := range eventBus.getStreamEvents() { + if payload.Data != nil && payload.Data.Type == streams.EventTypeComplete { + if payload.Data.TurnID != "turn-1" { + t.Fatalf("completion turn_id = %q, want turn-1", payload.Data.TurnID) + } + return + } + } + t.Fatal("no complete stream event published") +} + func TestHandleAgentEvent_RecordsStreamedAssistantTextForResumeContext(t *testing.T) { mgr, _ := createTestManagerWithTracking() history, err := NewSessionHistoryManager(t.TempDir(), "", newTestLogger()) diff --git a/apps/backend/internal/agent/runtime/lifecycle/manager_launch.go b/apps/backend/internal/agent/runtime/lifecycle/manager_launch.go index daa28e86ba..2d7bf2654b 100644 --- a/apps/backend/internal/agent/runtime/lifecycle/manager_launch.go +++ b/apps/backend/internal/agent/runtime/lifecycle/manager_launch.go @@ -571,6 +571,9 @@ func (m *Manager) launchPrepareRequest(req *LaunchRequest, profileInfo *AgentPro if req.SessionID != "" { reqWithWorktree.Metadata["session_id"] = req.SessionID } + if req.TurnID != "" { + reqWithWorktree.Metadata["prompt_turn_id"] = req.TurnID + } if err := mergeRouteOverrideEnv(&reqWithWorktree); err != nil { return LaunchRequest{}, "", err @@ -725,6 +728,7 @@ func (m *Manager) launchBuildExecutorRequest(ctx context.Context, executionID st TaskEnvironmentID: reqWithWorktree.TaskEnvironmentID, AgentProfileID: executionProfileID(reqWithWorktree), OfficeAgentProfileID: reqWithWorktree.AgentProfileID, + PromptTurnID: reqWithWorktree.TurnID, WorkspacePath: reqWithWorktree.WorkspacePath, WorkspaceSourceRoots: workspaceSourceRoots(reqWithWorktree.WorkspaceFolders, workspaceRepositorySpecsFromLaunch(reqWithWorktree)), Protocol: string(agentConfig.Runtime().Protocol), @@ -1596,6 +1600,18 @@ func (m *Manager) SetExecutionDescription(_ context.Context, executionID string, return nil } +// SetPromptTurnID binds the next prompt completion to a durable Kandev turn. +// The value is kept on the in-memory execution so it can be snapshotted onto +// the terminal stream event before AgentReady can admit a successor prompt. +func (m *Manager) SetPromptTurnID(_ context.Context, executionID, turnID string) error { + execution, exists := m.executionStore.Get(executionID) + if !exists { + return fmt.Errorf("execution %q not found", executionID) + } + execution.setPromptTurnID(turnID) + return nil +} + // SetExecutionEnv stores per-run environment variables for the next agent subprocess start. func (m *Manager) SetExecutionEnv(_ context.Context, executionID string, env map[string]string) error { execution, exists := m.executionStore.Get(executionID) diff --git a/apps/backend/internal/agent/runtime/lifecycle/types.go b/apps/backend/internal/agent/runtime/lifecycle/types.go index 7929cf01fb..2067871399 100644 --- a/apps/backend/internal/agent/runtime/lifecycle/types.go +++ b/apps/backend/internal/agent/runtime/lifecycle/types.go @@ -78,7 +78,12 @@ type AgentExecution struct { // out a generation that was merely admitted (not yet dispatched — its buffers // are still being reset) or already completed. dispatchedPromptGeneration uint64 - promptLifecycleMu sync.Mutex + // promptTurnID is the durable Kandev turn bound to the currently + // dispatched prompt. It is snapshotted onto terminal AgentEvents before + // AgentReady is published, so a queued successor cannot overwrite the + // completion's attribution while its stream frame is in flight. + promptTurnID string + promptLifecycleMu sync.Mutex // PrepareResult carries the environment preparation result back to the caller // so it can be persisted synchronously before UpdateTaskSession clobbers metadata. @@ -303,6 +308,24 @@ func (e *AgentExecution) promptGenerationSnapshot() uint64 { return e.promptGeneration } +func (e *AgentExecution) promptTurnIDSnapshot() string { + if e == nil { + return "" + } + e.promptLifecycleMu.Lock() + defer e.promptLifecycleMu.Unlock() + return e.promptTurnID +} + +func (e *AgentExecution) setPromptTurnID(turnID string) { + if e == nil { + return + } + e.promptLifecycleMu.Lock() + e.promptTurnID = turnID + e.promptLifecycleMu.Unlock() +} + // GetAgentCtlClient returns the agentctl client for this execution func (ae *AgentExecution) GetAgentCtlClient() *agentctl.Client { return ae.agentctl @@ -656,6 +679,7 @@ type LaunchRequest struct { // backward-compatible behavior by using AgentProfileID. ExecutionProfileID string StartAgent bool // Transfer launch activity through initial startup/prompt + TurnID string // Durable Kandev turn for the initial prompt, when present WorkspacePath string // Host path to workspace (original repository path) TaskDescription string // Task description to send via ACP prompt Attachments []MessageAttachment // Attachments (images/files) for the initial prompt diff --git a/apps/backend/internal/agentctl/server/adapter/transport/acp/adapter_prompt.go b/apps/backend/internal/agentctl/server/adapter/transport/acp/adapter_prompt.go index a9bc64a6be..c662cb8019 100644 --- a/apps/backend/internal/agentctl/server/adapter/transport/acp/adapter_prompt.go +++ b/apps/backend/internal/agentctl/server/adapter/transport/acp/adapter_prompt.go @@ -272,16 +272,17 @@ func (a *Adapter) sendPrompt( // subscriber sees an approximate input count. It has no input/output split, // so Estimated remains true. usage_update cost is cumulative session cost; // consumeUsageDelta converts it to the current turn's nonnegative delta. - delta, costSubcents := a.consumeUsageDelta(sessionID) + delta, costSubcents, costPresent := a.consumeUsageDeltaWithPresence(sessionID) if usage == nil { - if delta > 0 || costSubcents > 0 { + if delta > 0 || costPresent { usage = &streams.PromptUsage{ InputTokens: delta, Estimated: true, ProviderReportedCostSubcents: costSubcents, + ProviderReportedCostPresent: costPresent, } } - } else if costSubcents > 0 { + } else if costPresent { // claude-acp: usage_update.cost.amount carries authoritative cumulative // USD cost — attach the derived turn delta so Layer A wins // downstream and the office cost subscriber stores the row @@ -289,6 +290,7 @@ func (a *Adapter) sendPrompt( // model id is a logical alias (sonnet / haiku) that won't match // any pricing entry, so this is the only accurate cost path. usage.ProviderReportedCostSubcents = costSubcents + usage.ProviderReportedCostPresent = true } a.sendUpdate(AgentEvent{ Type: streams.EventTypeComplete, diff --git a/apps/backend/internal/agentctl/server/adapter/transport/acp/adapter_updates.go b/apps/backend/internal/agentctl/server/adapter/transport/acp/adapter_updates.go index d9273912e4..eb93269bd4 100644 --- a/apps/backend/internal/agentctl/server/adapter/transport/acp/adapter_updates.go +++ b/apps/backend/internal/agentctl/server/adapter/transport/acp/adapter_updates.go @@ -397,6 +397,7 @@ type usageTracker struct { consumedUsed int64 latestCostSubcents int64 consumedCostSubcents int64 + costObserved bool // maxSize is the largest context-window size reported for this session. // claude-acp's default model emits a 200K turn-start frame and a 1M // turn-end frame; sticky-max prevents the stale 200K from shrinking @@ -434,6 +435,7 @@ func (a *Adapter) retainConsumedUsageBaselineLocked(sessionID string) { a.usageBySession[sessionID] = tracker tracker.consumedUsed = tracker.latestUsed tracker.consumedCostSubcents = tracker.latestCostSubcents + tracker.costObserved = false } // recordUsageAndMaxSize updates per-session usage/cost tracking and returns @@ -453,6 +455,7 @@ func (a *Adapter) recordUsageAndMaxSize(sessionID string, size, used int64, cost } tr.latestUsed = used if cost != nil && strings.EqualFold(cost.Currency, "USD") { + tr.costObserved = true costSubcents := max(int64(cost.Amount*10000), 0) if costSubcents < tr.latestCostSubcents { tr.consumedCostSubcents = costSubcents @@ -477,17 +480,29 @@ func (a *Adapter) resetContextWindowMaxSize(sessionID string) { // previous prompt boundary and the delta in cumulative USD session cost. Both // consumed baselines advance to the latest observed sample. func (a *Adapter) consumeUsageDelta(sessionID string) (int64, int64) { + delta, cost, _ := a.consumeUsageDeltaWithPresence(sessionID) + return delta, cost +} + +// consumeUsageDeltaWithPresence returns nonnegative growth in context +// occupancy, the cumulative USD cost delta, and whether the provider supplied +// a USD cost sample for this turn. The presence bit is intentionally separate +// from cost: an explicit zero is authoritative and must not fall through to +// list-price estimation. +func (a *Adapter) consumeUsageDeltaWithPresence(sessionID string) (int64, int64, bool) { a.mu.Lock() defer a.mu.Unlock() tr := a.usageBySession[sessionID] if tr == nil { - return 0, 0 + return 0, 0, false } delta := max(tr.latestUsed-tr.consumedUsed, 0) cost := max(tr.latestCostSubcents-tr.consumedCostSubcents, 0) + present := tr.costObserved tr.consumedUsed = tr.latestUsed tr.consumedCostSubcents = tr.latestCostSubcents - return delta, cost + tr.costObserved = false + return delta, cost, present } // consumeUsageBaselineLocked marks replayed usage and cost as historical so @@ -497,6 +512,7 @@ func (a *Adapter) consumeUsageBaselineLocked(sessionID string) { if tr := a.usageBySession[sessionID]; tr != nil { tr.consumedUsed = tr.latestUsed tr.consumedCostSubcents = tr.latestCostSubcents + tr.costObserved = false } } diff --git a/apps/backend/internal/agentctl/server/adapter/transport/acp/usage_test.go b/apps/backend/internal/agentctl/server/adapter/transport/acp/usage_test.go index e39399f54f..d020b2a5fb 100644 --- a/apps/backend/internal/agentctl/server/adapter/transport/acp/usage_test.go +++ b/apps/backend/internal/agentctl/server/adapter/transport/acp/usage_test.go @@ -229,6 +229,25 @@ func TestUsageTracker_CumulativeUSDCost(t *testing.T) { } } +func TestUsageTracker_ExplicitZeroCostIsPresent(t *testing.T) { + a := newTestAdapter() + const sess = "sess-zero-cost" + + a.convertUsageUpdate(sess, usageUpdateWithCost(200_000, 100, 0, "USD")) + _, cost, present := a.consumeUsageDeltaWithPresence(sess) + if cost != 0 { + t.Fatalf("cost = %d, want 0", cost) + } + if !present { + t.Fatal("zero USD cost should be reported as present") + } + + _, _, present = a.consumeUsageDeltaWithPresence(sess) + if present { + t.Fatal("consumed cost presence should not leak into the next turn") + } +} + func TestUsageTracker_IgnoresNonUSDCostAndResetsDecreases(t *testing.T) { a := newTestAdapter() const sess = "sess-cost-currency" diff --git a/apps/backend/internal/agentctl/types/streams/agent.go b/apps/backend/internal/agentctl/types/streams/agent.go index 317e2e5674..2929ddcdd9 100644 --- a/apps/backend/internal/agentctl/types/streams/agent.go +++ b/apps/backend/internal/agentctl/types/streams/agent.go @@ -120,6 +120,11 @@ type AgentEvent struct { // attributed to a newer prompt on the same execution. PromptGeneration uint64 `json:"prompt_generation,omitempty"` + // TurnID is the durable Kandev turn identity captured for terminal events. + // It travels with the completion frame so consumers do not have to infer + // ownership from cross-subject event ordering. + TurnID string `json:"turn_id,omitempty"` + // --- Message fields (for "message_chunk" type) --- // Text contains streaming text content from the agent. @@ -505,7 +510,12 @@ type PromptUsage struct { ThoughtTokens int64 `json:"thought_tokens,omitempty"` TotalTokens int64 `json:"total_tokens"` ProviderReportedCostSubcents int64 `json:"provider_reported_cost_subcents,omitempty"` - Estimated bool `json:"estimated,omitempty"` + // ProviderReportedCostPresent distinguishes an explicit provider-reported + // zero from a missing cost sample. Providers can legitimately report zero + // for BYOK/free turns, and that value must still suppress list-price + // estimation downstream. + ProviderReportedCostPresent bool `json:"provider_reported_cost_present,omitempty"` + Estimated bool `json:"estimated,omitempty"` } // ToolCallContentItem represents a content item produced by a tool call. diff --git a/apps/backend/internal/backendapp/adapters.go b/apps/backend/internal/backendapp/adapters.go index 8577938cee..ff73b03f21 100644 --- a/apps/backend/internal/backendapp/adapters.go +++ b/apps/backend/internal/backendapp/adapters.go @@ -192,6 +192,7 @@ func buildLifecycleLaunchRequest( AgentProfileID: officeProfileID, ExecutionProfileID: req.AgentProfileID, StartAgent: req.StartAgent, + TurnID: req.TurnID, WorkspacePath: workspacePath, TaskDescription: req.TaskDescription, Attachments: convertToLifecycleAttachments(req.Attachments), @@ -320,6 +321,10 @@ func (a *lifecycleAdapter) SetExecutionDescription(ctx context.Context, agentExe return a.mgr.SetExecutionDescription(ctx, agentExecutionID, description) } +func (a *lifecycleAdapter) SetPromptTurnID(ctx context.Context, agentExecutionID, turnID string) error { + return a.mgr.SetPromptTurnID(ctx, agentExecutionID, turnID) +} + // RequiresCloneURL implements executor.ExecutorTypeCapabilities by delegating to // the lifecycle manager. Without this, executor types like local_docker and // sprites can't tell the orchestrator they need a clone URL. diff --git a/apps/backend/internal/office/costs/modelsdev/client.go b/apps/backend/internal/office/costs/modelsdev/client.go index 8528b48301..fb4214dcc4 100644 --- a/apps/backend/internal/office/costs/modelsdev/client.go +++ b/apps/backend/internal/office/costs/modelsdev/client.go @@ -51,7 +51,14 @@ type Client struct { index map[string]shared.ModelPricing info map[string]ModelInfo loadedAt time.Time - cacheBuf []byte // raw on-disk JSON (parsed lazily on miss) + // catalogGen increments on every catalogue install (warmFromDisk or + // refreshPhysical). CatalogVersion()'s RFC3339 string only has + // one-second resolution, so two installs landing in the same + // wall-clock second would report the same version — cacheIfVersionCurrent + // uses this counter instead, so it can't mistake a fresher install for + // the one a caller snapshotted. + catalogGen uint64 + cacheBuf []byte // raw on-disk JSON (parsed lazily on miss) } // ModelInfo holds non-pricing metadata from models.dev for a model. @@ -149,6 +156,97 @@ func (c *Client) LookupModelInfo(ctx context.Context, modelID string) (ModelInfo return ModelInfo{}, false } +// CatalogVersion implements shared.PricingCatalogVersioner. Returns the +// on-disk cache's load time in RFC3339 (UTC), or "" when nothing has been +// loaded yet (cold cache, no Lookup call has run). models.dev's dataset +// carries no version field of its own (see datasetEntry below) — the +// cache load/fetch time is the honest "as-of" identifier available. +func (c *Client) CatalogVersion() string { + c.mu.RLock() + defer c.mu.RUnlock() + return c.catalogVersionLocked() +} + +// catalogVersionLocked returns CatalogVersion's value assuming c.mu is +// already held (read or write) by the caller. Factored out so +// LookupForModelWithVersion can read pricing and version together under one +// lock acquisition instead of two independent ones. +func (c *Client) catalogVersionLocked() string { + if c.loadedAt.IsZero() { + return "" + } + return c.loadedAt.UTC().Format(time.RFC3339) +} + +// LookupForModelWithVersion implements shared.PricingLookupWithVersion. +// Identical lookup behavior to LookupForModel, but reads pricing and +// CatalogVersion from the same snapshot in every branch — including the +// cold-cache-buffer parse path, where the buffer and its version are +// captured together before parsing — so a background refresh landing +// mid-call can never pair one catalogue's rates with a different +// catalogue's version identifier (docs/specs/office/costs.md). +func (c *Client) LookupForModelWithVersion(ctx context.Context, modelID string) (shared.ModelPricing, string, bool) { + key, strategy := Normalize(modelID) + if strategy != StrategyLookup { + return shared.ModelPricing{}, "", false + } + c.once.Do(func() { c.warmFromDisk(ctx) }) + + c.mu.RLock() + pricing, ok := c.index[key] + version := c.catalogVersionLocked() + c.mu.RUnlock() + if ok { + c.maybeRefresh(ctx) + return pricing, version, true + } + + buf, bufVersion, bufGen := c.snapshotBufferAndVersion() + if len(buf) > 0 { + if pricing, ok = lookupInDataset(buf, key); ok { + c.cacheIfVersionCurrent(key, pricing, bufGen) + c.maybeRefresh(ctx) + return pricing, bufVersion, true + } + } + + c.maybeRefresh(ctx) + return shared.ModelPricing{}, "", false +} + +// cacheIfVersionCurrent stores pricing into the index under key, but only if +// the catalogue is still the one snapshotGen was captured from. Without this +// guard, a Refresh landing between the caller's snapshotBufferAndVersion call +// and this write would let a stale rate get written into the NEW index — +// refreshPhysical rebuilds c.index from the new buffer and replaces the map +// wholesale under c.mu — so a later, unrelated lookup for the same key would +// then read old rates paired with the new catalogue version, the exact +// provenance lie LookupForModelWithVersion exists to prevent +// (docs/specs/office/costs.md). Compares catalogGen rather than the RFC3339 +// version string: the string only has one-second resolution, so two installs +// landing in the same wall-clock second would otherwise compare equal and +// defeat this guard. The returned (pricing, bufVersion) pair for THIS call is +// unaffected either way, since both were derived from the same buf snapshot. +func (c *Client) cacheIfVersionCurrent(key string, pricing shared.ModelPricing, snapshotGen uint64) { + c.mu.Lock() + defer c.mu.Unlock() + if c.catalogGen != snapshotGen { + return + } + c.index[key] = pricing +} + +// snapshotBufferAndVersion returns the cache buffer, its catalogue version +// string, and its generation counter together under one lock, so a caller +// that parses buf afterward can report the version that actually produced +// it (and guard a delayed write with the generation) rather than whatever +// is current by the time the parse finishes. +func (c *Client) snapshotBufferAndVersion() ([]byte, string, uint64) { + c.mu.RLock() + defer c.mu.RUnlock() + return c.cacheBuf, c.catalogVersionLocked(), c.catalogGen +} + // warmFromDisk reads the cache file into cacheBuf so subsequent // lookups can parse individual model entries lazily. Missing or // unreadable cache is non-fatal — the next refresh tick warms it. @@ -176,6 +274,7 @@ func (c *Client) warmFromDisk(ctx context.Context) { c.mu.Lock() c.cacheBuf = buf c.loadedAt = stat.ModTime() + c.catalogGen++ c.mu.Unlock() if time.Since(stat.ModTime()) >= c.ttl { c.startBackgroundRefresh(context.WithoutCancel(ctx)) @@ -310,6 +409,7 @@ func (c *Client) refreshPhysical(ctx context.Context) error { c.index = newIndex c.info = newInfo c.loadedAt = time.Now().UTC() + c.catalogGen++ c.mu.Unlock() return nil } diff --git a/apps/backend/internal/office/costs/modelsdev/client_internal_test.go b/apps/backend/internal/office/costs/modelsdev/client_internal_test.go new file mode 100644 index 0000000000..512aac682e --- /dev/null +++ b/apps/backend/internal/office/costs/modelsdev/client_internal_test.go @@ -0,0 +1,72 @@ +package modelsdev + +import ( + "testing" + "time" + + "github.com/kandev/kandev/internal/office/shared" +) + +// TestClient_CacheIfVersionCurrent_SkipsStaleWrite is the deterministic +// regression test for the P1 cold-buffer-parse race: a caller that +// snapshotted an old catalogue generation must not write into the index once +// a concurrent Refresh has moved the catalogue on, or a later, unrelated +// lookup for the same key would read old rates paired with the new +// catalogue version. Runs entirely single-threaded (no goroutines, no +// timing-dependent race) by driving cacheIfVersionCurrent directly with a +// snapshot generation that is deliberately stale. +func TestClient_CacheIfVersionCurrent_SkipsStaleWrite(t *testing.T) { + c := &Client{index: make(map[string]shared.ModelPricing)} + c.loadedAt = time.Now().UTC() + c.catalogGen = 1 + staleGen := c.catalogGen + + // Simulate the concurrent Refresh that moved the catalogue on after the + // caller's snapshot was taken. + c.loadedAt = c.loadedAt.Add(2 * time.Second) + c.catalogGen++ + currentGen := c.catalogGen + + stalePricing := shared.ModelPricing{InputPerMillion: 999} + c.cacheIfVersionCurrent("stale-key", stalePricing, staleGen) + if _, ok := c.index["stale-key"]; ok { + t.Error("cacheIfVersionCurrent must not write when snapshotGen is stale (catalogue moved on)") + } + + freshPricing := shared.ModelPricing{InputPerMillion: 111} + c.cacheIfVersionCurrent("fresh-key", freshPricing, currentGen) + if got, ok := c.index["fresh-key"]; !ok || got != freshPricing { + t.Errorf("cacheIfVersionCurrent must write when snapshotGen matches the current catalogue, got %+v ok=%v", got, ok) + } +} + +// TestClient_CacheIfVersionCurrent_SameSecondInstallsAreDistinguishable is +// the regression test for the CodeRabbit finding on PR #2606: CatalogVersion +// is an RFC3339 string with one-second resolution, so two catalogue installs +// landing in the same wall-clock second would report the identical version +// string and defeat a string-based staleness guard. catalogGen has no such +// resolution limit — this asserts a stale snapshot is still rejected even +// when loadedAt (and therefore CatalogVersion()) does not advance at all. +func TestClient_CacheIfVersionCurrent_SameSecondInstallsAreDistinguishable(t *testing.T) { + c := &Client{index: make(map[string]shared.ModelPricing)} + now := time.Now().UTC() + c.loadedAt = now + c.catalogGen = 1 + staleGen := c.catalogGen + staleVersion := c.CatalogVersion() + + // Second install lands in the exact same wall-clock second: loadedAt + // (and thus CatalogVersion()'s RFC3339 string) is unchanged, but this is + // a genuinely distinct catalogue install. + c.loadedAt = now + c.catalogGen++ + if c.CatalogVersion() != staleVersion { + t.Fatal("test setup bug: CatalogVersion must collide within the same second") + } + + stalePricing := shared.ModelPricing{InputPerMillion: 999} + c.cacheIfVersionCurrent("stale-key", stalePricing, staleGen) + if _, ok := c.index["stale-key"]; ok { + t.Error("cacheIfVersionCurrent must reject a stale generation even when CatalogVersion()'s string is unchanged") + } +} diff --git a/apps/backend/internal/office/costs/modelsdev/client_test.go b/apps/backend/internal/office/costs/modelsdev/client_test.go index b3f1c53666..135ee4090d 100644 --- a/apps/backend/internal/office/costs/modelsdev/client_test.go +++ b/apps/backend/internal/office/costs/modelsdev/client_test.go @@ -13,6 +13,7 @@ import ( "github.com/kandev/kandev/internal/common/logger" "github.com/kandev/kandev/internal/office/costs/modelsdev" + "github.com/kandev/kandev/internal/office/shared" ) // sampleDataset mimics the models.dev /api.json shape: provider keys @@ -40,6 +41,18 @@ const sampleDataset = `{ } }` +// altDataset prices claude-opus-4-7 differently from sampleDataset (20/100 +// vs 15/75 USD per million input/output). Used only by the +// LookupForModelWithVersion concurrency test to make a mismatched +// (pricing, version) pairing observable. +const altDataset = `{ + "anthropic": { + "models": { + "claude-opus-4-7": {"cost": {"input": 20.0, "output": 100.0, "cache_read": 2.0, "cache_write": 25.0}} + } + } +}` + func newStubServer(t *testing.T, body string) *httptest.Server { t.Helper() srv := httptest.NewServer(http.HandlerFunc(func(w http.ResponseWriter, _ *http.Request) { @@ -116,6 +129,202 @@ func TestClient_Lookup(t *testing.T) { } } +// CatalogVersion implements shared.PricingCatalogVersioner: "" before +// anything has loaded, and a non-empty RFC3339 timestamp once the cache +// has warmed from disk or refreshed from the network — models.dev's +// dataset carries no version field of its own, so the writer needs this +// "as-of" signal to attribute a models_dev_list-priced row. +func TestClient_CatalogVersion(t *testing.T) { + dir := t.TempDir() + cachePath := filepath.Join(dir, "models-dev.json") + c, _ := newTestClient(t, cachePath) + + if v := c.CatalogVersion(); v != "" { + t.Fatalf("CatalogVersion before any load = %q, want empty", v) + } + + if err := c.Refresh(context.Background()); err != nil { + t.Fatalf("Refresh: %v", err) + } + + v := c.CatalogVersion() + if v == "" { + t.Fatal("CatalogVersion after refresh = empty, want a timestamp") + } + if _, err := time.Parse(time.RFC3339, v); err != nil { + t.Errorf("CatalogVersion = %q, not RFC3339: %v", v, err) + } +} + +// LookupForModelWithVersion implements shared.PricingLookupWithVersion: +// pricing matches LookupForModel and version matches CatalogVersion for the +// same warm cache state, and both are zero-valued together on a miss. +func TestClient_LookupForModelWithVersion(t *testing.T) { + dir := t.TempDir() + cachePath := filepath.Join(dir, "models-dev.json") + c, _ := newTestClient(t, cachePath) + + if err := c.Refresh(context.Background()); err != nil { + t.Fatalf("Refresh: %v", err) + } + + wantPricing, ok := c.LookupForModel(context.Background(), "claude-opus-4-7") + if !ok { + t.Fatal("expected hit on claude-opus-4-7") + } + wantVersion := c.CatalogVersion() + + pricing, version, ok := c.LookupForModelWithVersion(context.Background(), "claude-opus-4-7") + if !ok { + t.Fatal("expected hit on claude-opus-4-7") + } + if pricing != wantPricing { + t.Errorf("LookupForModelWithVersion pricing = %+v, want %+v", pricing, wantPricing) + } + if version != wantVersion { + t.Errorf("LookupForModelWithVersion version = %q, want %q", version, wantVersion) + } + + if pricing, version, ok := c.LookupForModelWithVersion(context.Background(), "claude-unknown-99"); ok || pricing != (shared.ModelPricing{}) || version != "" { + t.Errorf("expected miss on unknown model, got pricing=%+v version=%q ok=%v", pricing, version, ok) + } +} + +// TestClient_LookupForModelWithVersion_ConsistentUnderConcurrentRefresh is +// the regression test for the P1 "cost provenance lies" race: reading +// pricing and CatalogVersion via two independent lock acquisitions (as +// resolveCostForUsage did before this fix) lets a concurrent Refresh land +// in between and pair one catalogue's rates with a different catalogue's +// version identifier on the stored row. LookupForModelWithVersion must +// return the two from one atomic snapshot, so every observed pairing is one +// of the two catalogues actually served — never a hybrid. +// +// CatalogVersion has one-second resolution (it's RFC3339, by design — see +// its doc comment), so the base and alt refreshes are separated by a real +// sleep to guarantee distinguishable version strings; a tight refresh loop +// within the same wall-clock second would produce identical version labels +// for genuinely different catalogues and make this test unable to tell them +// apart. Reader goroutines hammer LookupForModelWithVersion continuously +// through the single base-to-alt transition, which is where the race +// window this test targets actually lives. +func TestClient_LookupForModelWithVersion_ConsistentUnderConcurrentRefresh(t *testing.T) { + dir := t.TempDir() + cachePath := filepath.Join(dir, "models-dev.json") + + var useAltDataset atomic.Bool + srv := httptest.NewServer(http.HandlerFunc(func(w http.ResponseWriter, _ *http.Request) { + if useAltDataset.Load() { + _, _ = w.Write([]byte(altDataset)) + return + } + _, _ = w.Write([]byte(sampleDataset)) + })) + t.Cleanup(srv.Close) + + log := logger.Default() + c := modelsdev.New(modelsdev.Config{ + CachePath: cachePath, + URL: srv.URL, + TTL: time.Hour, + HTTPClient: srv.Client(), + }, log) + + ctx := context.Background() + if err := c.Refresh(ctx); err != nil { + t.Fatalf("initial refresh: %v", err) + } + basePricing, baseVersion, ok := c.LookupForModelWithVersion(ctx, "claude-opus-4-7") + if !ok { + t.Fatal("expected hit on claude-opus-4-7 after initial refresh") + } + + // Guarantee the alt refresh lands in a different RFC3339 second than the + // base refresh above, so the two version strings are distinguishable. + time.Sleep(1100 * time.Millisecond) + + stop := make(chan struct{}) + var mismatches atomic.Int64 + + // altPricing/altVersion are set once by the main goroutine after the alt + // refresh below, and read concurrently by the reader goroutines started + // just above it; guarded by altMu (not plain vars) so that sharing is + // race-detector-clean, independent of the LookupForModelWithVersion + // atomicity this test is actually exercising. + var altMu sync.Mutex + var altPricing shared.ModelPricing + var altVersion string + getAlt := func() (shared.ModelPricing, string) { + altMu.Lock() + defer altMu.Unlock() + return altPricing, altVersion + } + + const readers = 8 + var readersWG sync.WaitGroup + readersWG.Add(readers) + for i := 0; i < readers; i++ { + go func() { + defer readersWG.Done() + for { + select { + case <-stop: + return + default: + } + pricing, version, ok := c.LookupForModelWithVersion(ctx, "claude-opus-4-7") + if !ok { + continue + } + wantAltPricing, wantAltVersion := getAlt() + switch version { + case baseVersion: + if pricing != basePricing { + mismatches.Add(1) + } + case wantAltVersion: + if wantAltVersion != "" && pricing != wantAltPricing { + mismatches.Add(1) + } + default: + // Neither known version yet (e.g. read raced ahead of the + // alt refresh completing) — not itself a mismatch, skip. + } + } + }() + } + + // Perform the single base->alt transition while readers are hammering + // the lookup concurrently — this Refresh call is exactly the race + // window the pre-fix two-lock pattern could straddle. + useAltDataset.Store(true) + if err := c.Refresh(ctx); err != nil { + t.Fatalf("alt refresh: %v", err) + } + newAltPricing, newAltVersion, ok := c.LookupForModelWithVersion(ctx, "claude-opus-4-7") + if !ok { + t.Fatal("expected hit on claude-opus-4-7 after alt refresh") + } + if newAltPricing == basePricing { + t.Fatal("test setup bug: altDataset must price claude-opus-4-7 differently from sampleDataset") + } + if newAltVersion == baseVersion { + t.Fatal("test setup bug: the alt refresh must produce a different catalogue version") + } + altMu.Lock() + altPricing, altVersion = newAltPricing, newAltVersion + altMu.Unlock() + + // Let readers keep hammering briefly against the now-settled alt state + // before stopping, so the alt-version branch above gets real coverage. + time.Sleep(50 * time.Millisecond) + close(stop) + readersWG.Wait() + + if mismatches.Load() > 0 { + t.Fatalf("observed %d mismatched (pricing, version) pairings — LookupForModelWithVersion is not atomic", mismatches.Load()) + } +} + // codex-acp model ids carry a / suffix and use dotted // versions. Normalize strips the effort; the dataset uses dotted form // too so the verbatim lookup hits. diff --git a/apps/backend/internal/office/models/enums.go b/apps/backend/internal/office/models/enums.go index f344a0deef..d5c55e5e5d 100644 --- a/apps/backend/internal/office/models/enums.go +++ b/apps/backend/internal/office/models/enums.go @@ -130,6 +130,37 @@ func (s BudgetScopeType) Valid() bool { return false } +// CostSource records which layer priced a CostEvent's dollars, distinct +// from Estimated (a token-synthesis flag). NULL on legacy rows written +// before this field existed — never inferred. +type CostSource string + +// Cost-source values. See docs/specs/office/costs.md. +const ( + // CostSourceProviderReported means the CLI forwarded an authoritative + // USD amount (claude-acp's usage_update.cost.amount); no rates apply. + CostSourceProviderReported CostSource = "provider_reported" + // CostSourceModelsDevList means the row was priced from a models.dev + // list-price lookup; the four rate_*_per_million columns and + // pricing_catalog_version record what was applied. + CostSourceModelsDevList CostSource = "models_dev_list" + // CostSourceUnpriced means no cost could be resolved (no pricing + // lookup wired, empty model id, or a lookup miss); cost_subcents is 0. + CostSourceUnpriced CostSource = "unpriced" +) + +// String implements fmt.Stringer. +func (s CostSource) String() string { return string(s) } + +// Valid reports whether s is one of the declared CostSource values. +func (s CostSource) Valid() bool { + switch s { + case CostSourceProviderReported, CostSourceModelsDevList, CostSourceUnpriced: + return true + } + return false +} + // BudgetPeriod is the time window a BudgetPolicy limit applies to. type BudgetPeriod string diff --git a/apps/backend/internal/office/models/models.go b/apps/backend/internal/office/models/models.go index 06b64f51d8..5b30421fe3 100644 --- a/apps/backend/internal/office/models/models.go +++ b/apps/backend/internal/office/models/models.go @@ -293,21 +293,42 @@ var ValidProjectStatuses = map[ProjectStatus]bool{ // is true when token counts were synthesised by the adapter (e.g. // cumulative-delta inference for codex-acp) rather than reported directly // by the agent; the row still counts toward budget totals at face value. +// +// TokensCachedRead / TokensCachedWrite / TurnID / UsageEventID / CostSource / +// the Rate*PerMillion columns / PricingCatalogVersion / CostContractVersion +// are all nullable (pointer fields): NULL means "not recorded" — a legacy +// row written before the column existed, or (for the cache split +// specifically) an adapter with no per-turn usage frame, such as codex-acp's +// occupancy-growth synthesis (Estimated=true). NULL is never backfilled to +// 0; TokensCachedIn keeps its original read+write sum semantics on every +// row so existing consumers of that column are unaffected. See +// docs/specs/office/costs.md. type CostEvent struct { - ID string `json:"id" db:"id"` - SessionID string `json:"session_id" db:"session_id"` - TaskID string `json:"task_id" db:"task_id"` - AgentProfileID string `json:"agent_profile_id" db:"agent_profile_id"` - ProjectID string `json:"project_id" db:"project_id"` - Model string `json:"model" db:"model"` - Provider string `json:"provider" db:"provider"` - TokensIn int64 `json:"tokens_in" db:"tokens_in"` - TokensCachedIn int64 `json:"tokens_cached_in" db:"tokens_cached_in"` - TokensOut int64 `json:"tokens_out" db:"tokens_out"` - CostSubcents int64 `json:"cost_subcents" db:"cost_subcents"` - Estimated bool `json:"estimated" db:"estimated"` - OccurredAt time.Time `json:"occurred_at" db:"occurred_at"` - CreatedAt time.Time `json:"created_at" db:"created_at"` + ID string `json:"id" db:"id"` + SessionID string `json:"session_id" db:"session_id"` + TaskID string `json:"task_id" db:"task_id"` + AgentProfileID string `json:"agent_profile_id" db:"agent_profile_id"` + ProjectID string `json:"project_id" db:"project_id"` + Model string `json:"model" db:"model"` + Provider string `json:"provider" db:"provider"` + TokensIn int64 `json:"tokens_in" db:"tokens_in"` + TokensCachedIn int64 `json:"tokens_cached_in" db:"tokens_cached_in"` + TokensCachedRead *int64 `json:"tokens_cached_read,omitempty" db:"tokens_cached_read"` + TokensCachedWrite *int64 `json:"tokens_cached_write,omitempty" db:"tokens_cached_write"` + TokensOut int64 `json:"tokens_out" db:"tokens_out"` + CostSubcents int64 `json:"cost_subcents" db:"cost_subcents"` + Estimated bool `json:"estimated" db:"estimated"` + TurnID *string `json:"turn_id,omitempty" db:"turn_id"` + UsageEventID *string `json:"usage_event_id,omitempty" db:"usage_event_id"` + CostSource *CostSource `json:"cost_source,omitempty" db:"cost_source"` + RateInputPerMillion *int64 `json:"rate_input_per_million,omitempty" db:"rate_input_per_million"` + RateCachedReadPerMillion *int64 `json:"rate_cached_read_per_million,omitempty" db:"rate_cached_read_per_million"` + RateCachedWritePerMillion *int64 `json:"rate_cached_write_per_million,omitempty" db:"rate_cached_write_per_million"` + RateOutputPerMillion *int64 `json:"rate_output_per_million,omitempty" db:"rate_output_per_million"` + PricingCatalogVersion *string `json:"pricing_catalog_version,omitempty" db:"pricing_catalog_version"` + CostContractVersion *int64 `json:"cost_contract_version,omitempty" db:"cost_contract_version"` + OccurredAt time.Time `json:"occurred_at" db:"occurred_at"` + CreatedAt time.Time `json:"created_at" db:"created_at"` } // BudgetPolicy represents a budget limit policy. LimitSubcents is diff --git a/apps/backend/internal/office/repository/sqlite/base.go b/apps/backend/internal/office/repository/sqlite/base.go index 1f11c51065..be56a86de2 100644 --- a/apps/backend/internal/office/repository/sqlite/base.go +++ b/apps/backend/internal/office/repository/sqlite/base.go @@ -109,6 +109,17 @@ func (r *Repository) ExecRaw(ctx context.Context, query string, args ...interfac // package should keep using typed methods. func (r *Repository) ReaderDB() *sqlx.DB { return r.ro } +// BeginTx starts a transaction on the shared writer connection. Exposed so +// the office service layer can make CreateCostEventTx atomic with another +// package's write (e.g. task.Repository.IncrementTaskSessionUsageTx via +// shared.SessionUsageWriterTx) — safe because both repositories are +// constructed from the same underlying *sqlx.DB (internal/backendapp/storage.go), +// so a transaction begun here can execute statements against tables either +// package owns. +func (r *Repository) BeginTx(ctx context.Context) (*sqlx.Tx, error) { + return r.db.BeginTxx(ctx, nil) +} + // initSchema creates all office tables if they don't exist. func (r *Repository) initSchema() error { if err := r.createCoreTables(); err != nil { @@ -258,6 +269,16 @@ func (r *Repository) createCostTables() error { // 10000 when rendering dollars. The estimated flag is set when // token counts were synthesised (e.g. cumulative-delta inference for // codex-acp) rather than reported directly by the agent. + // + // tokens_cached_read / tokens_cached_write / turn_id / usage_event_id / + // cost_source / rate_*_per_million / pricing_catalog_version / + // cost_contract_version have deliberately NO DEFAULT: an absent default + // is what makes a fresh row's un-set columns NULL rather than 0, matching + // the ALTER-based migration in migrateCostEventContract (which the same + // columns must stay byte-identical to — see base_migrations.go). NULL + // means "not recorded" (legacy row, or an adapter with no per-turn usage + // frame); 0 would silently claim zero cache activity. See + // docs/specs/office/costs.md. _, err := r.db.Exec(` CREATE TABLE IF NOT EXISTS office_cost_events ( id TEXT PRIMARY KEY, @@ -269,9 +290,20 @@ func (r *Repository) createCostTables() error { provider TEXT DEFAULT '', tokens_in INTEGER DEFAULT 0, tokens_cached_in INTEGER DEFAULT 0, + tokens_cached_read INTEGER, + tokens_cached_write INTEGER, tokens_out INTEGER DEFAULT 0, cost_subcents INTEGER NOT NULL DEFAULT 0, estimated INTEGER NOT NULL DEFAULT 0, + turn_id TEXT, + usage_event_id TEXT, + cost_source TEXT, + rate_input_per_million INTEGER, + rate_cached_read_per_million INTEGER, + rate_cached_write_per_million INTEGER, + rate_output_per_million INTEGER, + pricing_catalog_version TEXT, + cost_contract_version INTEGER, occurred_at TIMESTAMP NOT NULL, created_at TIMESTAMP NOT NULL ); @@ -286,6 +318,10 @@ func (r *Repository) createCostTables() error { -- full table scan per task_sessions row - measured at 76.72s at a modest -- 4,000 sessions / 80,000 events versus 0.17s indexed. CREATE INDEX IF NOT EXISTS idx_office_cost_events_session_id ON office_cost_events(session_id); + -- uniq_office_cost_usage_event is created by migrateCostEventContract + -- (base_migrations.go), not here: schema init runs before migrations, + -- so indexing usage_event_id inline would crash a pre-migration + -- database that doesn't have the column yet. CREATE TABLE IF NOT EXISTS office_budget_policies ( id TEXT PRIMARY KEY, diff --git a/apps/backend/internal/office/repository/sqlite/base_migrations.go b/apps/backend/internal/office/repository/sqlite/base_migrations.go index dbe14e4978..2da1d4b9a7 100644 --- a/apps/backend/internal/office/repository/sqlite/base_migrations.go +++ b/apps/backend/internal/office/repository/sqlite/base_migrations.go @@ -22,6 +22,7 @@ func (r *Repository) runMigrations() { r.migrateSchedulerColumns() r.migrateFailureColumns() r.migrateRunPayloadIndexes() + r.migrateCostEventContract() if err := r.migrateTaskPriorityToText(); err != nil { // Surface to stderr; this stage runs from initSchema which doesn't // have a logger handle. The recreate is wrapped in a transaction @@ -37,6 +38,42 @@ func (r *Repository) runMigrations() { r.migrateProviderRouting() } +// migrateCostEventContract adds the cache read/write split, turn +// attribution, and cost-provenance columns to office_cost_events for +// databases created before this contract (docs/specs/office/costs.md). Every +// ALTER is nullable with no DEFAULT: a legacy row must read NULL, never 0, +// because a merged tokens_cached_in cannot be decomposed and an unversioned +// "0" would be indistinguishable from "no cache activity". Keep this column +// set byte-identical to the inline CREATE TABLE in createCostTables so a +// fresh database and a migrated one converge. +func (r *Repository) migrateCostEventContract() { + r.migrate.Apply("office_cost_events.tokens_cached_read", + `ALTER TABLE office_cost_events ADD COLUMN tokens_cached_read INTEGER`) + r.migrate.Apply("office_cost_events.tokens_cached_write", + `ALTER TABLE office_cost_events ADD COLUMN tokens_cached_write INTEGER`) + r.migrate.Apply("office_cost_events.turn_id", + `ALTER TABLE office_cost_events ADD COLUMN turn_id TEXT`) + r.migrate.Apply("office_cost_events.usage_event_id", + `ALTER TABLE office_cost_events ADD COLUMN usage_event_id TEXT`) + r.migrate.Apply("office_cost_events.cost_source", + `ALTER TABLE office_cost_events ADD COLUMN cost_source TEXT`) + r.migrate.Apply("office_cost_events.rate_input_per_million", + `ALTER TABLE office_cost_events ADD COLUMN rate_input_per_million INTEGER`) + r.migrate.Apply("office_cost_events.rate_cached_read_per_million", + `ALTER TABLE office_cost_events ADD COLUMN rate_cached_read_per_million INTEGER`) + r.migrate.Apply("office_cost_events.rate_cached_write_per_million", + `ALTER TABLE office_cost_events ADD COLUMN rate_cached_write_per_million INTEGER`) + r.migrate.Apply("office_cost_events.rate_output_per_million", + `ALTER TABLE office_cost_events ADD COLUMN rate_output_per_million INTEGER`) + r.migrate.Apply("office_cost_events.pricing_catalog_version", + `ALTER TABLE office_cost_events ADD COLUMN pricing_catalog_version TEXT`) + r.migrate.Apply("office_cost_events.cost_contract_version", + `ALTER TABLE office_cost_events ADD COLUMN cost_contract_version INTEGER`) + r.migrate.Apply("uniq_office_cost_usage_event", + `CREATE UNIQUE INDEX IF NOT EXISTS uniq_office_cost_usage_event + ON office_cost_events(usage_event_id) WHERE usage_event_id IS NOT NULL`) +} + func (r *Repository) migrateRunPayloadIndexes() { r.migrate.Apply( "runs.idx_run_payload_comment_id", diff --git a/apps/backend/internal/office/repository/sqlite/cost_event_contract_migration_test.go b/apps/backend/internal/office/repository/sqlite/cost_event_contract_migration_test.go new file mode 100644 index 0000000000..b8d7c7c000 --- /dev/null +++ b/apps/backend/internal/office/repository/sqlite/cost_event_contract_migration_test.go @@ -0,0 +1,157 @@ +package sqlite_test + +import ( + "testing" + "time" + + "github.com/jmoiron/sqlx" + _ "github.com/mattn/go-sqlite3" + + "github.com/kandev/kandev/internal/office/repository/sqlite" +) + +// TestCostEventContractMigration_LegacyRowsGetNullColumns covers the +// legacy-column contract from docs/specs/office/costs.md: an office_cost_events +// table that predates the cache-split / cost-provenance / turn-attribution +// columns gets them added as nullable, a pre-existing row is left with NULL +// (never 0) in every new column, the partial unique index on +// usage_event_id exists, and the migration replays cleanly on a second boot. +func TestCostEventContractMigration_LegacyRowsGetNullColumns(t *testing.T) { + db, err := sqlx.Open("sqlite3", ":memory:") + if err != nil { + t.Fatalf("open sqlite: %v", err) + } + db.SetMaxOpenConns(1) + t.Cleanup(func() { _ = db.Close() }) + + // Seed the pre-migration shape of office_cost_events (as it exists on + // main today, per base.go's createCostTables) and a legacy row, before + // the repo (and its migration) ever runs. + if _, err := db.Exec(` + CREATE TABLE office_cost_events ( + id TEXT PRIMARY KEY, + session_id TEXT DEFAULT '', + task_id TEXT DEFAULT '', + agent_profile_id TEXT DEFAULT '', + project_id TEXT DEFAULT '', + model TEXT DEFAULT '', + provider TEXT DEFAULT '', + tokens_in INTEGER DEFAULT 0, + tokens_cached_in INTEGER DEFAULT 0, + tokens_out INTEGER DEFAULT 0, + cost_subcents INTEGER NOT NULL DEFAULT 0, + estimated INTEGER NOT NULL DEFAULT 0, + occurred_at TIMESTAMP NOT NULL, + created_at TIMESTAMP NOT NULL + )`); err != nil { + t.Fatalf("seed legacy office_cost_events: %v", err) + } + if _, err := db.Exec(` + INSERT INTO office_cost_events ( + id, task_id, tokens_in, tokens_cached_in, tokens_out, + cost_subcents, estimated, occurred_at, created_at + ) VALUES ( + 'legacy-1', 'task-legacy', 100, 40, 50, + 10, 0, CURRENT_TIMESTAMP, CURRENT_TIMESTAMP + )`); err != nil { + t.Fatalf("seed legacy row: %v", err) + } + + // First boot: initSchema's CREATE TABLE IF NOT EXISTS is a no-op + // (table already exists), so the new columns can only appear via the + // ALTER-based migration. + if _, err := sqlite.NewWithDB(db, db, nil); err != nil { + t.Fatalf("first boot: %v", err) + } + + newColumns := []string{ + "tokens_cached_read", + "tokens_cached_write", + "turn_id", + "usage_event_id", + "cost_source", + "rate_input_per_million", + "rate_cached_read_per_million", + "rate_cached_write_per_million", + "rate_output_per_million", + "pricing_catalog_version", + "cost_contract_version", + } + for _, col := range newColumns { + var got *string + if err := db.QueryRow(`SELECT ` + col + ` FROM office_cost_events WHERE id = 'legacy-1'`).Scan(&got); err != nil { + t.Fatalf("select %s on legacy row: %v", col, err) + } + if got != nil { + t.Errorf("legacy row %s = %v, want NULL (never 0)", col, *got) + } + } + + // tokens_cached_in itself must be untouched by the migration: it stays + // the pre-existing sum, not backfilled or reinterpreted. + var tokensCachedIn int64 + if err := db.QueryRow(`SELECT tokens_cached_in FROM office_cost_events WHERE id = 'legacy-1'`).Scan(&tokensCachedIn); err != nil { + t.Fatalf("select tokens_cached_in: %v", err) + } + if tokensCachedIn != 40 { + t.Errorf("tokens_cached_in = %d, want 40 (unchanged legacy sum)", tokensCachedIn) + } + + var indexName string + if err := db.Get(&indexName, ` + SELECT name FROM sqlite_master WHERE type = 'index' AND name = 'uniq_office_cost_usage_event' + `); err != nil { + t.Fatalf("uniq_office_cost_usage_event index is missing: %v", err) + } + + // Second boot: the migration must replay without error (idempotent + // ALTERs) and must not disturb the already-NULL legacy row. + if _, err := sqlite.NewWithDB(db, db, nil); err != nil { + t.Fatalf("second boot (replay): %v", err) + } + var tokensCachedInAfterReplay int64 + if err := db.QueryRow(`SELECT tokens_cached_in FROM office_cost_events WHERE id = 'legacy-1'`).Scan(&tokensCachedInAfterReplay); err != nil { + t.Fatalf("select tokens_cached_in after replay: %v", err) + } + if tokensCachedInAfterReplay != 40 { + t.Errorf("tokens_cached_in after replay = %d, want 40 (unchanged)", tokensCachedInAfterReplay) + } +} + +// TestCostEventContractMigration_UsageEventIDUniqueAllowsManyNulls confirms +// the partial unique index only constrains non-NULL usage_event_id values — +// unbounded legacy/manual rows with NULL usage_event_id must never collide. +func TestCostEventContractMigration_UsageEventIDUniqueAllowsManyNulls(t *testing.T) { + db, err := sqlx.Open("sqlite3", ":memory:") + if err != nil { + t.Fatalf("open sqlite: %v", err) + } + db.SetMaxOpenConns(1) + t.Cleanup(func() { _ = db.Close() }) + + if _, err := sqlite.NewWithDB(db, db, nil); err != nil { + t.Fatalf("init repo: %v", err) + } + + now := time.Now().UTC() + insertRow := func(id string, usageEventID interface{}) error { + _, err := db.Exec(` + INSERT INTO office_cost_events (id, cost_subcents, estimated, occurred_at, created_at, usage_event_id) + VALUES (?, 0, 0, ?, ?, ?) + `, id, now, now, usageEventID) + return err + } + + if err := insertRow("row-null-1", nil); err != nil { + t.Fatalf("insert row-null-1: %v", err) + } + if err := insertRow("row-null-2", nil); err != nil { + t.Fatalf("insert row-null-2 (NULL usage_event_id must not collide): %v", err) + } + if err := insertRow("row-dup-1", "usage-evt-dup"); err != nil { + t.Fatalf("insert row-dup-1: %v", err) + } + if err := insertRow("row-dup-2", "usage-evt-dup"); err == nil { + t.Fatal("expected unique violation on duplicate non-NULL usage_event_id, got nil error") + } +} diff --git a/apps/backend/internal/office/repository/sqlite/cost_event_contract_postgres_test.go b/apps/backend/internal/office/repository/sqlite/cost_event_contract_postgres_test.go new file mode 100644 index 0000000000..7993660dde --- /dev/null +++ b/apps/backend/internal/office/repository/sqlite/cost_event_contract_postgres_test.go @@ -0,0 +1,95 @@ +package sqlite_test + +import ( + "context" + "testing" + + "github.com/kandev/kandev/internal/office/repository/sqlite" + taskrepo "github.com/kandev/kandev/internal/task/repository/sqlite" + "github.com/kandev/kandev/internal/testutil" +) + +// TestPostgresCostEventContractMigration is the PostgreSQL twin of +// TestCostEventContractMigration_LegacyRowsGetNullColumns +// (docs/specs/office/costs.md, ADR 0027): a pre-migration office_cost_events +// table gets the cache-split / provenance / turn-attribution columns added +// nullable, a legacy row is left NULL (never 0), the partial unique index +// exists, and the migration replays cleanly on a second boot. Skips unless +// KANDEV_TEST_POSTGRES_DSN is set. +// +// The test inserts a cost event row that references a task_id, so the tasks +// table must exist first — initialize via taskrepo, mirroring production boot +// order (see internal/office/repository/sqlite/workflow_test.go). +func TestPostgresCostEventContractMigration(t *testing.T) { + db := testutil.OpenIsolatedPostgres(t, testutil.PostgresDSNFromEnv(t)) + ctx := context.Background() + + if _, err := taskrepo.NewWithDB(db, db, nil); err != nil { + t.Fatalf("init task repo: %v", err) + } + + if _, err := db.ExecContext(ctx, ` + CREATE TABLE office_cost_events ( + id TEXT PRIMARY KEY, + session_id TEXT DEFAULT '', + task_id TEXT DEFAULT '', + agent_profile_id TEXT DEFAULT '', + project_id TEXT DEFAULT '', + model TEXT DEFAULT '', + provider TEXT DEFAULT '', + tokens_in BIGINT DEFAULT 0, + tokens_cached_in BIGINT DEFAULT 0, + tokens_out BIGINT DEFAULT 0, + cost_subcents BIGINT NOT NULL DEFAULT 0, + estimated BOOLEAN NOT NULL DEFAULT false, + occurred_at TIMESTAMPTZ NOT NULL, + created_at TIMESTAMPTZ NOT NULL + )`); err != nil { + t.Fatalf("seed legacy office_cost_events: %v", err) + } + if _, err := db.ExecContext(ctx, ` + INSERT INTO office_cost_events ( + id, task_id, tokens_in, tokens_cached_in, tokens_out, + cost_subcents, estimated, occurred_at, created_at + ) VALUES ( + 'legacy-1', 'task-legacy', 100, 40, 50, + 10, false, now(), now() + )`); err != nil { + t.Fatalf("seed legacy row: %v", err) + } + + if _, err := sqlite.NewWithDB(db, db, nil); err != nil { + t.Fatalf("first boot: %v", err) + } + + newColumns := []string{ + "tokens_cached_read", "tokens_cached_write", "turn_id", "usage_event_id", + "cost_source", "rate_input_per_million", "rate_cached_read_per_million", + "rate_cached_write_per_million", "rate_output_per_million", + "pricing_catalog_version", "cost_contract_version", + } + for _, col := range newColumns { + var got *string + if err := db.QueryRowContext(ctx, + `SELECT `+col+`::text FROM office_cost_events WHERE id = 'legacy-1'`, + ).Scan(&got); err != nil { + t.Fatalf("select %s on legacy row: %v", col, err) + } + if got != nil { + t.Errorf("legacy row %s = %v, want NULL (never 0)", col, *got) + } + } + + var indexName string + if err := db.QueryRowContext(ctx, ` + SELECT indexname FROM pg_indexes + WHERE tablename = 'office_cost_events' AND indexname = 'uniq_office_cost_usage_event' + `).Scan(&indexName); err != nil { + t.Fatalf("uniq_office_cost_usage_event index is missing: %v", err) + } + + // Second boot: the migration must replay without error. + if _, err := sqlite.NewWithDB(db, db, nil); err != nil { + t.Fatalf("second boot (replay): %v", err) + } +} diff --git a/apps/backend/internal/office/repository/sqlite/cost_event_tx_atomicity_test.go b/apps/backend/internal/office/repository/sqlite/cost_event_tx_atomicity_test.go new file mode 100644 index 0000000000..2b184663ac --- /dev/null +++ b/apps/backend/internal/office/repository/sqlite/cost_event_tx_atomicity_test.go @@ -0,0 +1,144 @@ +package sqlite_test + +import ( + "context" + "testing" + "time" + + "github.com/jmoiron/sqlx" + _ "github.com/mattn/go-sqlite3" + + "github.com/kandev/kandev/internal/office/models" + "github.com/kandev/kandev/internal/office/repository/sqlite" + taskrepo "github.com/kandev/kandev/internal/task/repository/sqlite" +) + +// TestCostEventTxAtomicity_RealReposShareTransaction is the PR #2606 review +// round 1 follow-up (F2): the only prior coverage of +// recordCostEventAndRollup's shared-transaction path +// (event_subscribers.go) used a test double that discarded its *sqlx.Tx +// parameter, so neither session.go's `tx != nil` branch nor +// CreateCostEventTx executing against a real transaction was ever exercised. +// +// This builds the real office and task repositories over one *sqlx.DB (the +// NewWithDB(db, db, nil) pattern already used by +// cost_event_contract_migration_test.go and workflow_test.go — safe because +// office.Repository.BeginTx and task.Repository.IncrementTaskSessionUsageTx +// are constructed from the same underlying writer connection in production, +// see BeginTx's doc comment), then drives BeginTx -> CreateCostEventTx -> +// IncrementTaskSessionUsageTx -> Rollback and asserts both office_cost_events +// and task_sessions are unchanged, followed by a committed retry that +// updates both. +func TestCostEventTxAtomicity_RealReposShareTransaction(t *testing.T) { + db, err := sqlx.Open("sqlite3", ":memory:") + if err != nil { + t.Fatalf("open sqlite: %v", err) + } + db.SetMaxOpenConns(1) + t.Cleanup(func() { _ = db.Close() }) + + taskRepo, err := taskrepo.NewWithDB(db, db, nil) + if err != nil { + t.Fatalf("init task repo: %v", err) + } + officeRepo, err := sqlite.NewWithDB(db, db, nil) + if err != nil { + t.Fatalf("init office repo: %v", err) + } + + ctx := context.Background() + now := time.Now().UTC() + if _, err := db.ExecContext(ctx, ` + INSERT INTO workspaces (id, name, created_at, updated_at) VALUES ('ws-atomic', 'ws', ?, ?) + `, now, now); err != nil { + t.Fatalf("seed workspace: %v", err) + } + if _, err := db.ExecContext(ctx, ` + INSERT INTO tasks (id, workspace_id, title, created_at, updated_at) VALUES ('task-atomic', 'ws-atomic', 't', ?, ?) + `, now, now); err != nil { + t.Fatalf("seed task: %v", err) + } + if _, err := db.ExecContext(ctx, ` + INSERT INTO task_sessions (id, task_id, started_at, updated_at) VALUES ('session-atomic', 'task-atomic', ?, ?) + `, now, now); err != nil { + t.Fatalf("seed session: %v", err) + } + + countCostEvents := func() int { + t.Helper() + var n int + if err := db.GetContext(ctx, &n, `SELECT COUNT(*) FROM office_cost_events`); err != nil { + t.Fatalf("count cost events: %v", err) + } + return n + } + rollupTotals := func() (tokensIn, tokensCachedIn, tokensOut, costSubcents int64) { + t.Helper() + row := db.QueryRowContext(ctx, ` + SELECT tokens_in, tokens_cached_in, tokens_out, cost_subcents + FROM task_sessions WHERE id = 'session-atomic' + `) + if err := row.Scan(&tokensIn, &tokensCachedIn, &tokensOut, &costSubcents); err != nil { + t.Fatalf("read rollup totals: %v", err) + } + return tokensIn, tokensCachedIn, tokensOut, costSubcents + } + buildEvent := func(usageEventID string) *models.CostEvent { + return &models.CostEvent{ + SessionID: "session-atomic", + TaskID: "task-atomic", + CostSubcents: 500, + UsageEventID: &usageEventID, + OccurredAt: now, + } + } + + // Rollback: both writes land in the same transaction; rolling it back + // must undo both, not just the one whose caller happened to fail. + tx, err := officeRepo.BeginTx(ctx) + if err != nil { + t.Fatalf("begin tx: %v", err) + } + if err := officeRepo.CreateCostEventTx(ctx, tx, buildEvent("usage-evt-rollback")); err != nil { + t.Fatalf("create cost event in tx: %v", err) + } + if err := taskRepo.IncrementTaskSessionUsageTx(ctx, tx, "session-atomic", 100, 40, 50, 500); err != nil { + t.Fatalf("increment rollup in tx: %v", err) + } + if err := tx.Rollback(); err != nil { + t.Fatalf("rollback: %v", err) + } + + if n := countCostEvents(); n != 0 { + t.Fatalf("cost events after rollback = %d, want 0", n) + } + if tokensIn, tokensCachedIn, tokensOut, costSubcents := rollupTotals(); tokensIn != 0 || tokensCachedIn != 0 || tokensOut != 0 || costSubcents != 0 { + t.Fatalf("rollup totals after rollback = (in=%d cached=%d out=%d cost=%d), want all zero", + tokensIn, tokensCachedIn, tokensOut, costSubcents) + } + + // Commit: a fresh transaction with the same shape commits both writes + // together, proving the office-begun transaction can actually reach + // task_sessions, a table the task package owns. + tx2, err := officeRepo.BeginTx(ctx) + if err != nil { + t.Fatalf("begin tx (commit path): %v", err) + } + if err := officeRepo.CreateCostEventTx(ctx, tx2, buildEvent("usage-evt-commit")); err != nil { + t.Fatalf("create cost event in tx (commit path): %v", err) + } + if err := taskRepo.IncrementTaskSessionUsageTx(ctx, tx2, "session-atomic", 100, 40, 50, 500); err != nil { + t.Fatalf("increment rollup in tx (commit path): %v", err) + } + if err := tx2.Commit(); err != nil { + t.Fatalf("commit: %v", err) + } + + if n := countCostEvents(); n != 1 { + t.Fatalf("cost events after commit = %d, want 1", n) + } + if tokensIn, tokensCachedIn, tokensOut, costSubcents := rollupTotals(); tokensIn != 100 || tokensCachedIn != 40 || tokensOut != 50 || costSubcents != 500 { + t.Fatalf("rollup totals after commit = (in=%d cached=%d out=%d cost=%d), want (100, 40, 50, 500)", + tokensIn, tokensCachedIn, tokensOut, costSubcents) + } +} diff --git a/apps/backend/internal/office/repository/sqlite/costs.go b/apps/backend/internal/office/repository/sqlite/costs.go index 9f5052a7ad..9786a0b0b9 100644 --- a/apps/backend/internal/office/repository/sqlite/costs.go +++ b/apps/backend/internal/office/repository/sqlite/costs.go @@ -2,30 +2,118 @@ package sqlite import ( "context" + "database/sql" + "errors" + "strings" "time" "github.com/google/uuid" + "github.com/jackc/pgx/v5/pgconn" + "github.com/jmoiron/sqlx" "github.com/kandev/kandev/internal/office/models" ) -// CreateCostEvent records a new cost event. +// ErrDuplicateUsageEvent is returned by CreateCostEvent when the row's +// UsageEventID collides with an existing one. The real duplicate source is +// not bus redelivery — neither the memory nor the NATS event bus redelivers +// (see internal/events/bus/{memory,nats}.go) — it is the SAME underlying +// completion frame reaching the publisher twice, e.g. a reconnecting WS +// client replaying a buffered stream event (see +// orchestrator.usageEventIDFor's doc comment for how the id is derived to +// collide in exactly that case). Either way this must not double-count +// cost; callers treat this as an idempotent no-op, not a failure, but still +// get a distinct signal to attribute to writer health rather than to +// "written". +var ErrDuplicateUsageEvent = errors.New("duplicate usage_event_id") + +// usageEventIndexName is the partial unique index enforcing at most one row +// per non-NULL usage_event_id (docs/specs/office/costs.md). +const usageEventIndexName = "uniq_office_cost_usage_event" + +// sqliteUsageEventViolationMessage is the substring go-sqlite3 puts in a +// UNIQUE-constraint error for this index. SQLite exposes no typed access to +// the violated index's name, only the column ("UNIQUE constraint failed: +// office_cost_events.usage_event_id"), which is unique to this index in +// this table. +const sqliteUsageEventViolationMessage = "UNIQUE constraint failed: office_cost_events.usage_event_id" + +// isUsageEventUniqueViolation reports whether err is a violation of +// uniq_office_cost_usage_event specifically, not any unique violation. On +// PostgreSQL it inspects the typed pgconn.PgError's constraint name; on +// SQLite (no typed access to the constraint name) it matches the +// column-list message documented above. Mirrors +// internal/task/repository/sqlite/task_external_id.go's +// isExternalIDUniqueViolation. +func isUsageEventUniqueViolation(err error) bool { + if err == nil { + return false + } + var pgErr *pgconn.PgError + if errors.As(err, &pgErr) { + return pgErr.Code == "23505" && pgErr.ConstraintName == usageEventIndexName + } + return strings.Contains(err.Error(), sqliteUsageEventViolationMessage) +} + +// CreateCostEvent records a new cost event. A UsageEventID collision +// (redelivery of the same prompt-usage event) is reported as +// ErrDuplicateUsageEvent rather than the raw driver error, so callers can +// treat it as an idempotent no-op. +// +// Delegates to CreateCostEventTx using r.db as the executor; a caller that +// needs this atomic with another write (e.g. the office cost subscriber's +// session-usage rollup) should call the Tx variant directly with a shared +// transaction instead. func (r *Repository) CreateCostEvent(ctx context.Context, event *models.CostEvent) error { + return r.CreateCostEventTx(ctx, nil, event) +} + +// CreateCostEventTx is CreateCostEvent's transactional twin: executes +// against tx when non-nil (falling back to r.db, the shared writer +// connection, when tx is nil) so a caller can make this atomic with another +// write in the same transaction — see BeginTx and +// shared.SessionUsageWriterTx's doc comment for why (docs/specs/office/costs.md, +// PR #2606 review). +func (r *Repository) CreateCostEventTx(ctx context.Context, tx *sqlx.Tx, event *models.CostEvent) error { if event.ID == "" { event.ID = uuid.New().String() } event.CreatedAt = time.Now().UTC() - _, err := r.db.ExecContext(ctx, r.db.Rebind(` + var exec interface { + ExecContext(ctx context.Context, query string, args ...interface{}) (sql.Result, error) + Rebind(query string) string + } + if tx != nil { + exec = tx + } else { + exec = r.db + } + + _, err := exec.ExecContext(ctx, exec.Rebind(` INSERT INTO office_cost_events ( id, session_id, task_id, agent_profile_id, project_id, - model, provider, tokens_in, tokens_cached_in, tokens_out, - cost_subcents, estimated, occurred_at, created_at - ) VALUES (?, ?, ?, ?, ?, ?, ?, ?, ?, ?, ?, ?, ?, ?) + model, provider, tokens_in, tokens_cached_in, + tokens_cached_read, tokens_cached_write, tokens_out, + cost_subcents, estimated, turn_id, usage_event_id, cost_source, + rate_input_per_million, rate_cached_read_per_million, + rate_cached_write_per_million, rate_output_per_million, + pricing_catalog_version, cost_contract_version, + occurred_at, created_at + ) VALUES (?, ?, ?, ?, ?, ?, ?, ?, ?, ?, ?, ?, ?, ?, ?, ?, ?, ?, ?, ?, ?, ?, ?, ?, ?) `), event.ID, event.SessionID, event.TaskID, event.AgentProfileID, event.ProjectID, event.Model, event.Provider, event.TokensIn, - event.TokensCachedIn, event.TokensOut, event.CostSubcents, - event.Estimated, event.OccurredAt, event.CreatedAt) + event.TokensCachedIn, event.TokensCachedRead, event.TokensCachedWrite, + event.TokensOut, event.CostSubcents, event.Estimated, + event.TurnID, event.UsageEventID, event.CostSource, + event.RateInputPerMillion, event.RateCachedReadPerMillion, + event.RateCachedWritePerMillion, event.RateOutputPerMillion, + event.PricingCatalogVersion, event.CostContractVersion, + event.OccurredAt, event.CreatedAt) + if err != nil && isUsageEventUniqueViolation(err) { + return ErrDuplicateUsageEvent + } return err } diff --git a/apps/backend/internal/office/repository/sqlite/costs_test.go b/apps/backend/internal/office/repository/sqlite/costs_test.go index 41de5fa575..c95cfff204 100644 --- a/apps/backend/internal/office/repository/sqlite/costs_test.go +++ b/apps/backend/internal/office/repository/sqlite/costs_test.go @@ -2,6 +2,7 @@ package sqlite_test import ( "context" + "errors" "testing" "time" @@ -87,6 +88,56 @@ func TestCostEvent_CreateAndList(t *testing.T) { } } +// TestCreateCostEvent_DuplicateUsageEventID confirms redelivery of the same +// prompt-usage bus event (same UsageEventID) is reported as +// ErrDuplicateUsageEvent rather than inserting a second row or a raw driver +// error, and that unrelated rows with a NULL UsageEventID never collide. +func TestCreateCostEvent_DuplicateUsageEventID(t *testing.T) { + repo := newTestRepo(t) + ctx := context.Background() + seedTasksTable(t, repo, "task-dup", "ws-dup") + + usageEventID := "usage-evt-1" + first := &models.CostEvent{ + TaskID: "task-dup", + CostSubcents: 10, + UsageEventID: &usageEventID, + OccurredAt: time.Now().UTC(), + } + if err := repo.CreateCostEvent(ctx, first); err != nil { + t.Fatalf("create first: %v", err) + } + + second := &models.CostEvent{ + TaskID: "task-dup", + CostSubcents: 10, + UsageEventID: &usageEventID, + OccurredAt: time.Now().UTC(), + } + err := repo.CreateCostEvent(ctx, second) + if !errors.Is(err, sqlite.ErrDuplicateUsageEvent) { + t.Fatalf("create second (duplicate) err = %v, want ErrDuplicateUsageEvent", err) + } + + costs, err := repo.ListCostEvents(ctx, "ws-dup") + if err != nil { + t.Fatalf("list costs: %v", err) + } + if len(costs) != 1 { + t.Fatalf("cost count = %d, want 1 (duplicate must not insert)", len(costs)) + } + + // Two rows with no UsageEventID (manual entries, or rows predating this + // field) must not collide with each other. + for i := 0; i < 2; i++ { + if err := repo.CreateCostEvent(ctx, &models.CostEvent{ + TaskID: "task-dup", CostSubcents: 5, OccurredAt: time.Now().UTC(), + }); err != nil { + t.Fatalf("create nil-usage-event row %d: %v", i, err) + } + } +} + func TestCostBreakdowns(t *testing.T) { repo := newTestRepo(t) ctx := context.Background() diff --git a/apps/backend/internal/office/service/cost_metrics.go b/apps/backend/internal/office/service/cost_metrics.go new file mode 100644 index 0000000000..c203c51b3d --- /dev/null +++ b/apps/backend/internal/office/service/cost_metrics.go @@ -0,0 +1,71 @@ +package service + +import ( + "expvar" + "strings" + + "go.uber.org/zap" +) + +// expvar maps published at package init, exposed via stdlib's /debug/vars +// handler (dev mode). Mirrors internal/office/scheduler/metrics_vars.go's +// shape: counters only, "k1=v1;k2=v2" label keys so a future Prometheus +// translation layer can split without re-shaping storage. +// +// This pair exists because handlePromptUsage previously had two silent +// early returns (decode failure, task lookup failure) and a silent insert +// failure — "if this writer silently stops, nothing today would notice" +// (docs/specs/office/costs.md). +var ( + costEventsWrittenTotal = expvar.NewMap("cost_events_written_total") + costEventsDroppedTotal = expvar.NewMap("cost_events_dropped_total") +) + +// Reasons a prompt-usage event never became a stored cost_events row. +const ( + costDropReasonDecodeError = "decode_error" + costDropReasonMissingIDs = "missing_ids" + costDropReasonTaskFieldsError = "task_fields_error" + costDropReasonInsertError = "insert_error" + costDropReasonDuplicate = "duplicate" +) + +const ( + metricCostEventWritten = "cost.metric.event_written" + metricCostEventDropped = "cost.metric.event_dropped" +) + +func costMetricLabel(pairs ...string) string { + if len(pairs)%2 != 0 { + return "" + } + parts := make([]string, 0, len(pairs)/2) + for i := 0; i < len(pairs); i += 2 { + parts = append(parts, pairs[i]+"="+pairs[i+1]) + } + return strings.Join(parts, ";") +} + +// recordCostEventWritten emits a structured log AND bumps the expvar +// counter under cost_events_written_total when handlePromptUsage +// successfully inserts a row. +func (s *Service) recordCostEventWritten(source, provider string) { + costEventsWrittenTotal.Add(costMetricLabel("source", source, "provider", provider), 1) + if s.logger == nil { + return + } + s.logger.Debug(metricCostEventWritten, + zap.String("cost_source", source), zap.String("provider", provider)) +} + +// recordCostEventDropped emits a structured log AND bumps the expvar +// counter under cost_events_dropped_total whenever a prompt-usage event +// does not result in a stored row. +func (s *Service) recordCostEventDropped(reason, taskID string) { + costEventsDroppedTotal.Add(costMetricLabel("reason", reason), 1) + if s.logger == nil { + return + } + s.logger.Warn(metricCostEventDropped, + zap.String("reason", reason), zap.String("task_id", taskID)) +} diff --git a/apps/backend/internal/office/service/event_subscribers.go b/apps/backend/internal/office/service/event_subscribers.go index 8b5b84a586..2450c4976c 100644 --- a/apps/backend/internal/office/service/event_subscribers.go +++ b/apps/backend/internal/office/service/event_subscribers.go @@ -6,7 +6,6 @@ import ( "encoding/json" "errors" "fmt" - "time" "go.uber.org/zap" @@ -117,21 +116,26 @@ type AgentLifecycleData struct { } type PromptUsageData struct { - TaskID string `json:"task_id"` - SessionID string `json:"session_id"` - AgentID string `json:"agent_id"` - AgentType string `json:"agent_type"` - Model string `json:"model"` - Provider string `json:"provider"` - Usage UsageTokens `json:"usage"` + TaskID string `json:"task_id"` + SessionID string `json:"session_id"` + AgentID string `json:"agent_id"` + AgentType string `json:"agent_type"` + Model string `json:"model"` + Provider string `json:"provider"` + Usage UsageTokens `json:"usage"` + TurnID string `json:"turn_id,omitempty"` + UsageEventID string `json:"usage_event_id,omitempty"` } // UsageTokens mirrors streams.PromptUsage on the wire. All counts are int64 // to match the stream type and to handle workspaces that accumulate over a // million tokens. ProviderReportedCostSubcents is forwarded from claude-acp's -// usage_update.cost.amount (USD float * 10000); when > 0 the subscriber uses -// it directly and skips the models.dev lookup. Estimated is true when the -// adapter synthesised tokens (codex-acp cumulative-delta inference). +// usage_update.cost.amount (USD float * 10000); when +// ProviderReportedCostPresent is true (including an explicit zero), the +// subscriber uses it directly and skips the models.dev lookup. The legacy +// positive-value check remains in the resolver for older events. Estimated is +// true when the adapter synthesised tokens (codex-acp cumulative-delta +// inference). type UsageTokens struct { InputTokens int64 `json:"input_tokens"` OutputTokens int64 `json:"output_tokens"` @@ -140,6 +144,7 @@ type UsageTokens struct { ThoughtTokens int64 `json:"thought_tokens,omitempty"` TotalTokens int64 `json:"total_tokens,omitempty"` ProviderReportedCostSubcents int64 `json:"provider_reported_cost_subcents,omitempty"` + ProviderReportedCostPresent bool `json:"provider_reported_cost_present,omitempty"` Estimated bool `json:"estimated,omitempty"` } @@ -521,58 +526,65 @@ func (s *Service) tryPostStartFallback( } // handlePromptUsage records a cost event from a session/prompt usage -// update. Cost resolution follows the three-layer order from -// docs/specs/office-costs/spec.md: +// update. Cost resolution follows the two-layer order from +// docs/specs/office/costs.md, and CostSource on the row records which +// layer actually produced the dollar amount (see resolveCostForUsage in +// prompt_usage_cost.go — distinct from Estimated, a token-synthesis flag): // // 1. Provider-reported cost (Layer A) — claude-acp emits exact USD per // turn on usage_update.cost.amount; the adapter forwards this as -// ProviderReportedCostSubcents. When > 0 the row is recorded -// verbatim and pricing lookup is skipped. This is the only accurate +// ProviderReportedCostSubcents plus a presence bit. When present (even +// when zero), the row is recorded verbatim and pricing lookup is skipped. +// This is the only accurate // path for claude-acp, whose model identifiers are logical aliases // (default / sonnet / haiku) with no real-name mapping. // 2. models.dev (Layer B) — when tokens are reported but no cost, // normalize the model id and look up pricing. On miss the row -// records cost_subcents=0 with estimated=true. +// records cost_subcents=0 with cost_source=unpriced; Estimated tracks +// data.Usage.Estimated verbatim on this path, not hardcoded true. // -// After insert the session totals (tokens_in / tokens_cached_in / tokens_out / -// cost_subcents) are incremented on task_sessions, and any applicable budget -// policy is evaluated. Estimated rows count toward budget totals at face value. +// buildCostEvent (prompt_usage_cost.go) also records the cache read/write +// split (NULL when Usage.Estimated — see its doc comment) and turn_id / +// usage_event_id threaded from the publish site. The ledger insert and the +// task_sessions rollup increment (tokens_in / tokens_cached_in / tokens_out / +// cost_subcents) are written atomically by recordCostEventAndRollup — see its +// doc comment for why non-atomic writes here are a real defect, not a +// theoretical one — and any applicable budget policy is evaluated afterward. +// Estimated rows count toward budget totals at face value. Every early +// return and the insert path record a writer-health metric (cost_metrics.go) +// so a silently-stopped writer is observable. func (s *Service) handlePromptUsage(ctx context.Context, event *bus.Event) error { data, err := decodeEventData[PromptUsageData](event) - if err != nil || data.TaskID == "" || data.SessionID == "" { + if err != nil { + s.recordCostEventDropped(costDropReasonDecodeError, "") + return nil + } + if data.TaskID == "" || data.SessionID == "" { + s.recordCostEventDropped(costDropReasonMissingIDs, data.TaskID) return nil } fields, err := s.repo.GetTaskExecutionFields(ctx, data.TaskID) if err != nil { + s.recordCostEventDropped(costDropReasonTaskFieldsError, data.TaskID) return nil } - costSubcents, estimated := s.resolveCostForUsage(ctx, *data) + resolution := s.resolveCostForUsage(ctx, *data) provider := resolveProvider(*data) - tokensCachedIn := data.Usage.CachedReadTokens + data.Usage.CachedWriteTokens + costEvent := buildCostEvent(*data, fields, s.projectIDForTask(ctx, data.TaskID), provider, resolution) - costEvent := &models.CostEvent{ - SessionID: data.SessionID, - TaskID: data.TaskID, - AgentProfileID: fields.AssigneeAgentProfileID, - ProjectID: s.projectIDForTask(ctx, data.TaskID), - Model: data.Model, - Provider: provider, - TokensIn: data.Usage.InputTokens, - TokensCachedIn: tokensCachedIn, - TokensOut: data.Usage.OutputTokens, - CostSubcents: costSubcents, - Estimated: estimated, - OccurredAt: time.Now().UTC(), - } - if err := s.repo.CreateCostEvent(ctx, costEvent); err != nil { + if err := s.recordCostEventAndRollup( + ctx, costEvent, data.SessionID, + data.Usage.InputTokens, costEvent.TokensCachedIn, data.Usage.OutputTokens, resolution.costSubcents, + ); err != nil { + if errors.Is(err, sqlite.ErrDuplicateUsageEvent) { + s.recordCostEventDropped(costDropReasonDuplicate, data.TaskID) + return nil + } + s.recordCostEventDropped(costDropReasonInsertError, data.TaskID) return err } - - s.incrementSessionUsageTotals( - ctx, data.SessionID, - data.Usage.InputTokens, tokensCachedIn, data.Usage.OutputTokens, costSubcents, - ) + s.recordCostEventWritten(string(resolution.source), provider) if fields.WorkspaceID != "" { if err := s.CheckBudget( @@ -585,37 +597,59 @@ func (s *Service) handlePromptUsage(ctx context.Context, event *bus.Event) error return nil } -// resolveCostForUsage applies the Layer A / Layer B lookup. Returns -// (costSubcents, estimated). Layer A wins when the adapter forwarded a -// non-zero provider-reported cost (claude-acp's usage_update.cost.amount). -// Layer B (models.dev) is queried when a PricingLookup is wired; on miss -// or when no PricingLookup is configured the row records 0/estimated. -func (s *Service) resolveCostForUsage( - ctx context.Context, data PromptUsageData, -) (int64, bool) { - if data.Usage.ProviderReportedCostSubcents > 0 { - return data.Usage.ProviderReportedCostSubcents, data.Usage.Estimated - } - if s.pricingLookup == nil || data.Model == "" { - return 0, true - } - pricing, ok := s.pricingLookup.LookupForModel(ctx, data.Model) +// recordCostEventAndRollup inserts costEvent and increments the +// task_sessions rollup atomically when the wired sessionUsageWriter +// supports it (shared.SessionUsageWriterTx) — every production wiring does +// (backendapp's SetSessionUsageWriter(repos.Task)). +// +// Atomicity matters because of how this card's own usage_event_id +// idempotency guard interacts with a partial failure: without a shared +// transaction, a rollup-increment failure after a successful ledger insert +// was logged and swallowed, and a later redelivery of the same completion +// (see ErrDuplicateUsageEvent's doc comment for why redelivery is real, not +// hypothetical) would hit the unique index and be dropped as a duplicate +// before the rollup ever got a second chance — task_sessions permanently +// behind office_cost_events with no recovery path (PR #2606 review). Wrapping +// both writes in one transaction closes that: any failure, including the +// rollup increment, rolls back the ledger insert too, so nothing is +// committed and a redelivered completion with the same usage_event_id +// retries the whole operation cleanly instead of colliding with a +// half-landed row. +// +// Falls back to the pre-existing non-atomic two-step write when the wired +// writer doesn't implement SessionUsageWriterTx (e.g. a simplified test +// double that doesn't need transactional coupling, or sessionUsageWriter +// left unset). +func (s *Service) recordCostEventAndRollup( + ctx context.Context, costEvent *models.CostEvent, sessionID string, + tokensIn, tokensCachedIn, tokensOut, costSubcents int64, +) error { + txWriter, ok := s.sessionUsageWriter.(shared.SessionUsageWriterTx) if !ok { - return 0, true + if err := s.repo.CreateCostEvent(ctx, costEvent); err != nil { + return err + } + s.incrementSessionUsageTotals(ctx, sessionID, tokensIn, tokensCachedIn, tokensOut, costSubcents) + return nil } - cost := costs.CalculateCostSubcents( - data.Usage.InputTokens, - data.Usage.CachedReadTokens, - data.Usage.CachedWriteTokens, - data.Usage.OutputTokens, - costs.ModelPricing{ - InputPerMillion: pricing.InputPerMillion, - CachedReadPerMillion: pricing.CachedReadPerMillion, - CachedWritePerMillion: pricing.CachedWritePerMillion, - OutputPerMillion: pricing.OutputPerMillion, - }, - ) - return cost, data.Usage.Estimated + + tx, err := s.repo.BeginTx(ctx) + if err != nil { + return err + } + if err := s.repo.CreateCostEventTx(ctx, tx, costEvent); err != nil { + _ = tx.Rollback() + return err + } + if err := txWriter.IncrementTaskSessionUsageTx( + ctx, tx, sessionID, tokensIn, tokensCachedIn, tokensOut, costSubcents, + ); err != nil { + _ = tx.Rollback() + s.logger.Warn("increment task_session usage failed, rolled back cost event insert", + zap.String("session_id", sessionID), zap.Error(err)) + return err + } + return tx.Commit() } // resolveProvider derives the provider id for the cost row. AgentType diff --git a/apps/backend/internal/office/service/event_subscribers_cost_atomicity_test.go b/apps/backend/internal/office/service/event_subscribers_cost_atomicity_test.go new file mode 100644 index 0000000000..1bd48c92d1 --- /dev/null +++ b/apps/backend/internal/office/service/event_subscribers_cost_atomicity_test.go @@ -0,0 +1,136 @@ +package service_test + +import ( + "context" + "errors" + "sync" + "testing" + + "github.com/jmoiron/sqlx" + + "github.com/kandev/kandev/internal/events" + "github.com/kandev/kandev/internal/events/bus" +) + +// errInjectedRollupFailure is returned by txSessionUsageWriter while armed, +// standing in for a real rollup-side failure (e.g. a constraint violation or +// a dropped connection) that happens after the cost-event insert has +// already been staged in the same transaction. +var errInjectedRollupFailure = errors.New("injected rollup failure") + +// txSessionUsageWriter is a shared.SessionUsageWriterTx test double whose +// Tx-aware increment can be armed to fail a fixed number of times before +// succeeding. It exists to prove recordCostEventAndRollup +// (event_subscribers.go) rolls the cost-event insert back when the rollup +// increment fails inside the shared transaction, and that a redelivery of +// the same usage_event_id afterward is not treated as a duplicate — the +// recovery path the PR #2606 review asked to see covered, since neither +// half of the write had actually committed. +type txSessionUsageWriter struct { + mu sync.Mutex + failNextN int + calls int + + tokensIn int64 + tokensCachedIn int64 + tokensOut int64 + costSubcents int64 +} + +func (w *txSessionUsageWriter) IncrementTaskSessionUsage( + ctx context.Context, sessionID string, tokensIn, tokensCachedIn, tokensOut, costSubcents int64, +) error { + return w.IncrementTaskSessionUsageTx(ctx, nil, sessionID, tokensIn, tokensCachedIn, tokensOut, costSubcents) +} + +func (w *txSessionUsageWriter) IncrementTaskSessionUsageTx( + _ context.Context, _ *sqlx.Tx, _ string, tokensIn, tokensCachedIn, tokensOut, costSubcents int64, +) error { + w.mu.Lock() + defer w.mu.Unlock() + w.calls++ + if w.failNextN > 0 { + w.failNextN-- + return errInjectedRollupFailure + } + w.tokensIn += tokensIn + w.tokensCachedIn += tokensCachedIn + w.tokensOut += tokensOut + w.costSubcents += costSubcents + return nil +} + +// TestPromptUsage_RollupFailureRollsBackLedgerInsertThenRetrySucceeds covers +// the recovery path the PR #2606 review specifically asked for: force the +// rollup update to fail, then retry the same usage_event_id. Before +// recordCostEventAndRollup wrapped both writes in one transaction, the +// ledger insert committed regardless of the rollup outcome; a redelivery +// carrying the same usage_event_id would then hit the unique index and be +// dropped as a duplicate, leaving task_sessions permanently behind +// office_cost_events with no way to catch up. +func TestPromptUsage_RollupFailureRollsBackLedgerInsertThenRetrySucceeds(t *testing.T) { + svc, eb := newTestServiceWithBus(t) + ctx := context.Background() + writer := &txSessionUsageWriter{failNextN: 1} + svc.SetSessionUsageWriter(writer) + + createTestAgent(t, svc, "ws-1", "worker-atomic") + insertTestTask(t, svc, "task-atomic", "ws-1") + setTestTaskAssignee(t, svc, "task-atomic", "worker-atomic") + + buildEvent := func() *bus.Event { + return bus.NewEvent(events.SessionPromptUsageUpdated, "test", map[string]interface{}{ + "task_id": "task-atomic", + "session_id": "session-atomic", + "agent_id": "claude-acp", + "model": "default", + "usage_event_id": "usage-evt-atomic-1", + "usage": map[string]interface{}{ + "input_tokens": 100, + "output_tokens": 50, + }, + }) + } + subject := events.BuildSessionPromptUsageSubject("session-atomic") + + // First delivery: the rollup increment is armed to fail. If the ledger + // insert and the rollup increment were not atomic, the insert would + // already be committed here. + if err := eb.Publish(ctx, subject, buildEvent()); err != nil { + t.Fatalf("publish (failing rollup): %v", err) + } + + costs, err := svc.ListCostEvents(ctx, "ws-1") + if err != nil { + t.Fatalf("list costs after failed rollup: %v", err) + } + if len(costs) != 0 { + t.Fatalf("cost events after failed rollup = %d, want 0 (ledger insert must roll back with the failed rollup)", + len(costs)) + } + if writer.tokensIn != 0 || writer.tokensOut != 0 { + t.Fatalf("rollup totals after failed attempt = (in=%d out=%d), want (0, 0): a failed increment must not partially apply", + writer.tokensIn, writer.tokensOut) + } + + // Redelivery of the same usage_event_id, this time with the rollup + // writer succeeding. Because nothing committed on the first attempt, + // the unique index must not treat this as a duplicate. + if err := eb.Publish(ctx, subject, buildEvent()); err != nil { + t.Fatalf("publish (retry): %v", err) + } + + costs, err = svc.ListCostEvents(ctx, "ws-1") + if err != nil { + t.Fatalf("list costs after retry: %v", err) + } + if len(costs) != 1 { + t.Fatalf("cost events after retry = %d, want 1", len(costs)) + } + if writer.tokensIn != 100 || writer.tokensOut != 50 { + t.Errorf("rollup totals after retry = (in=%d out=%d), want (in=100 out=50)", writer.tokensIn, writer.tokensOut) + } + if writer.calls != 2 { + t.Errorf("rollup writer called %d times, want 2 (one failed attempt, one successful retry)", writer.calls) + } +} diff --git a/apps/backend/internal/office/service/event_subscribers_cost_test.go b/apps/backend/internal/office/service/event_subscribers_cost_test.go new file mode 100644 index 0000000000..b5cc6a87c3 --- /dev/null +++ b/apps/backend/internal/office/service/event_subscribers_cost_test.go @@ -0,0 +1,382 @@ +package service_test + +import ( + "context" + "testing" + + "github.com/kandev/kandev/internal/events" + "github.com/kandev/kandev/internal/events/bus" + "github.com/kandev/kandev/internal/office/models" + "github.com/kandev/kandev/internal/office/shared" +) + +// fakePricingLookup is a minimal shared.PricingLookup test double. hit +// controls whether LookupForModel reports a hit; version (if non-empty) +// makes the fake also satisfy shared.PricingCatalogVersioner, so tests can +// exercise both the with- and without-versioner cases. +type fakePricingLookup struct { + hit bool + pricing shared.ModelPricing + version string +} + +func (f *fakePricingLookup) LookupForModel(_ context.Context, _ string) (shared.ModelPricing, bool) { + if !f.hit { + return shared.ModelPricing{}, false + } + return f.pricing, true +} + +func (f *fakePricingLookup) CatalogVersion() string { return f.version } + +// TestPromptUsage_CacheSplitRecordedWhenNotEstimated confirms the cache +// read/write split reaches storage intact — the P1 defect was that +// tokens_cached_in = read + write was computed and then the split was +// thrown away one line before the INSERT, even though CalculateCostSubcents +// already prices read and write at distinct per-million rates. +func TestPromptUsage_CacheSplitRecordedWhenNotEstimated(t *testing.T) { + svc, eb := newTestServiceWithBus(t) + ctx := context.Background() + + createTestAgent(t, svc, "ws-1", "worker-split") + insertTestTask(t, svc, "task-split", "ws-1") + setTestTaskAssignee(t, svc, "task-split", "worker-split") + + event := bus.NewEvent(events.SessionPromptUsageUpdated, "test", map[string]interface{}{ + "task_id": "task-split", + "session_id": "session-split", + "agent_id": "claude-acp", + "model": "sonnet", + "usage": map[string]interface{}{ + "input_tokens": 100, + "cached_read_tokens": 40, + "cached_write_tokens": 60, + "output_tokens": 10, + }, + }) + if err := eb.Publish(ctx, events.BuildSessionPromptUsageSubject("session-split"), event); err != nil { + t.Fatalf("publish: %v", err) + } + + costs, err := svc.ListCostEvents(ctx, "ws-1") + if err != nil || len(costs) != 1 { + t.Fatalf("list costs: %v (len=%d)", err, len(costs)) + } + row := costs[0] + if row.TokensCachedRead == nil || *row.TokensCachedRead != 40 { + t.Errorf("TokensCachedRead = %v, want 40", row.TokensCachedRead) + } + if row.TokensCachedWrite == nil || *row.TokensCachedWrite != 60 { + t.Errorf("TokensCachedWrite = %v, want 60", row.TokensCachedWrite) + } + // tokens_cached_in keeps its original sum semantics for existing + // consumers (e.g. the tree-holds rollup and card 2faa29da). + if row.TokensCachedIn != 100 { + t.Errorf("TokensCachedIn = %d, want 100 (read+write, unchanged semantics)", row.TokensCachedIn) + } +} + +// TestPromptUsage_CacheSplitNullWhenEstimated confirms the codex/no-usage- +// frame caveat: when Usage.Estimated is true (context-occupancy synthesis, +// adapter_prompt.go's fallback), the split columns must be NULL, never 0 — +// a zero would falsely claim "no cache activity" for an adapter that never +// reported cache tokens at all. tokens_cached_in still equals read+write +// (both zero here), preserving the existing column's semantics. +func TestPromptUsage_CacheSplitNullWhenEstimated(t *testing.T) { + svc, eb := newTestServiceWithBus(t) + ctx := context.Background() + + createTestAgent(t, svc, "ws-1", "worker-codex-null") + insertTestTask(t, svc, "task-codex-null", "ws-1") + setTestTaskAssignee(t, svc, "task-codex-null", "worker-codex-null") + + event := bus.NewEvent(events.SessionPromptUsageUpdated, "test", map[string]interface{}{ + "task_id": "task-codex-null", + "session_id": "session-codex-null", + "agent_id": "codex-acp", + "model": "gpt-5.4-mini", + "usage": map[string]interface{}{ + "input_tokens": 350, + "estimated": true, + }, + }) + if err := eb.Publish(ctx, events.BuildSessionPromptUsageSubject("session-codex-null"), event); err != nil { + t.Fatalf("publish: %v", err) + } + + costs, err := svc.ListCostEvents(ctx, "ws-1") + if err != nil || len(costs) != 1 { + t.Fatalf("list costs: %v (len=%d)", err, len(costs)) + } + row := costs[0] + if row.TokensCachedRead != nil { + t.Errorf("TokensCachedRead = %v, want nil (NULL, never 0)", *row.TokensCachedRead) + } + if row.TokensCachedWrite != nil { + t.Errorf("TokensCachedWrite = %v, want nil (NULL, never 0)", *row.TokensCachedWrite) + } +} + +// TestPromptUsage_CostSourceProviderReported confirms Layer A rows are +// tagged provider_reported with no rates recorded (rates only apply to a +// models.dev list-price calculation). +func TestPromptUsage_CostSourceProviderReported(t *testing.T) { + svc, eb := newTestServiceWithBus(t) + ctx := context.Background() + + createTestAgent(t, svc, "ws-1", "worker-src-a") + insertTestTask(t, svc, "task-src-a", "ws-1") + setTestTaskAssignee(t, svc, "task-src-a", "worker-src-a") + + event := bus.NewEvent(events.SessionPromptUsageUpdated, "test", map[string]interface{}{ + "task_id": "task-src-a", + "session_id": "session-src-a", + "agent_id": "claude-acp", + "model": "sonnet", + "usage": map[string]interface{}{ + "input_tokens": 100, + "output_tokens": 200, + "provider_reported_cost_subcents": 616, + }, + }) + if err := eb.Publish(ctx, events.BuildSessionPromptUsageSubject("session-src-a"), event); err != nil { + t.Fatalf("publish: %v", err) + } + + costs, err := svc.ListCostEvents(ctx, "ws-1") + if err != nil || len(costs) != 1 { + t.Fatalf("list costs: %v (len=%d)", err, len(costs)) + } + row := costs[0] + if row.CostSource == nil || *row.CostSource != models.CostSourceProviderReported { + t.Fatalf("CostSource = %v, want %q", row.CostSource, models.CostSourceProviderReported) + } + if row.RateInputPerMillion != nil { + t.Errorf("RateInputPerMillion = %v, want nil on the provider-reported path", *row.RateInputPerMillion) + } + if row.PricingCatalogVersion != nil { + t.Errorf("PricingCatalogVersion = %v, want nil on the provider-reported path", *row.PricingCatalogVersion) + } +} + +func TestPromptUsage_ExplicitZeroProviderCostSkipsPricing(t *testing.T) { + svc, eb := newTestServiceWithBus(t) + ctx := context.Background() + svc.SetPricingLookup(&fakePricingLookup{ + hit: true, + pricing: shared.ModelPricing{InputPerMillion: 999999, OutputPerMillion: 999999}, + }) + + createTestAgent(t, svc, "ws-1", "worker-zero-provider") + insertTestTask(t, svc, "task-zero-provider", "ws-1") + setTestTaskAssignee(t, svc, "task-zero-provider", "worker-zero-provider") + + event := bus.NewEvent(events.SessionPromptUsageUpdated, "test", map[string]interface{}{ + "task_id": "task-zero-provider", + "session_id": "session-zero-provider", + "agent_id": "claude-acp", + "model": "sonnet", + "usage": map[string]interface{}{ + "input_tokens": 100, + "output_tokens": 200, + "provider_reported_cost_subcents": 0, + "provider_reported_cost_present": true, + }, + }) + if err := eb.Publish(ctx, events.BuildSessionPromptUsageSubject("session-zero-provider"), event); err != nil { + t.Fatalf("publish: %v", err) + } + + costs, err := svc.ListCostEvents(ctx, "ws-1") + if err != nil || len(costs) != 1 { + t.Fatalf("list costs: %v (len=%d)", err, len(costs)) + } + row := costs[0] + if row.CostSubcents != 0 { + t.Fatalf("cost_subcents = %d, want 0", row.CostSubcents) + } + if row.CostSource == nil || *row.CostSource != models.CostSourceProviderReported { + t.Fatalf("CostSource = %v, want %q", row.CostSource, models.CostSourceProviderReported) + } +} + +// TestPromptUsage_CostSourceModelsDevList confirms Layer B rows are tagged +// models_dev_list with the four applied rates and the catalogue version +// recorded — the whole point of the provenance field: today `estimated` +// cannot distinguish "priced from a provider amount" from "priced from a +// list-price calculation", which downstream read as "metered" for 671 of +// 673 events until corrected on the dashboard side. +func TestPromptUsage_CostSourceModelsDevList(t *testing.T) { + svc, eb := newTestServiceWithBus(t) + ctx := context.Background() + svc.SetPricingLookup(&fakePricingLookup{ + hit: true, + pricing: shared.ModelPricing{ + InputPerMillion: 300000, + CachedReadPerMillion: 30000, + CachedWritePerMillion: 375000, + OutputPerMillion: 1500000, + }, + version: "2026-08-12T00:00:00Z", + }) + + createTestAgent(t, svc, "ws-1", "worker-src-b") + insertTestTask(t, svc, "task-src-b", "ws-1") + setTestTaskAssignee(t, svc, "task-src-b", "worker-src-b") + + event := bus.NewEvent(events.SessionPromptUsageUpdated, "test", map[string]interface{}{ + "task_id": "task-src-b", + "session_id": "session-src-b", + "model": "claude-sonnet-4-5", + "usage": map[string]interface{}{ + "input_tokens": 1_000_000, + "output_tokens": 1_000_000, + }, + }) + if err := eb.Publish(ctx, events.BuildSessionPromptUsageSubject("session-src-b"), event); err != nil { + t.Fatalf("publish: %v", err) + } + + costs, err := svc.ListCostEvents(ctx, "ws-1") + if err != nil || len(costs) != 1 { + t.Fatalf("list costs: %v (len=%d)", err, len(costs)) + } + row := costs[0] + if row.CostSource == nil || *row.CostSource != models.CostSourceModelsDevList { + t.Fatalf("CostSource = %v, want %q", row.CostSource, models.CostSourceModelsDevList) + } + if row.RateInputPerMillion == nil || *row.RateInputPerMillion != 300000 { + t.Errorf("RateInputPerMillion = %v, want 300000", row.RateInputPerMillion) + } + if row.RateOutputPerMillion == nil || *row.RateOutputPerMillion != 1500000 { + t.Errorf("RateOutputPerMillion = %v, want 1500000", row.RateOutputPerMillion) + } + if row.PricingCatalogVersion == nil || *row.PricingCatalogVersion != "2026-08-12T00:00:00Z" { + t.Errorf("PricingCatalogVersion = %v, want the fake's version", row.PricingCatalogVersion) + } + if row.CostContractVersion == nil || *row.CostContractVersion != 2 { + t.Errorf("CostContractVersion = %v, want 2 (in-band activation point)", row.CostContractVersion) + } +} + +// TestPromptUsage_CostSourceUnpriced confirms the "both layers miss" case +// (no provider-reported cost, no pricing lookup wired) is tagged unpriced, +// not silently left as a bare estimated=true with no source at all. It also +// covers the R2-F4 regression: usage carrying authoritative (non-synthesised) +// token counts must keep Estimated=false even though the row is unpriced — +// cost_source=unpriced alone carries "we could not resolve a price"; +// Estimated is strictly "the token counts were synthesised" (see +// costContractVersion's v1→v2 doc comment in prompt_usage_cost.go). Before +// that fix this case incorrectly forced Estimated=true. +func TestPromptUsage_CostSourceUnpriced(t *testing.T) { + svc, eb := newTestServiceWithBus(t) + ctx := context.Background() + + createTestAgent(t, svc, "ws-1", "worker-src-c") + insertTestTask(t, svc, "task-src-c", "ws-1") + setTestTaskAssignee(t, svc, "task-src-c", "worker-src-c") + + event := bus.NewEvent(events.SessionPromptUsageUpdated, "test", map[string]interface{}{ + "task_id": "task-src-c", + "session_id": "session-src-c", + "model": "butler_a", + "usage": map[string]interface{}{ + "input_tokens": 100, + "output_tokens": 200, + "estimated": false, + }, + }) + if err := eb.Publish(ctx, events.BuildSessionPromptUsageSubject("session-src-c"), event); err != nil { + t.Fatalf("publish: %v", err) + } + + costs, err := svc.ListCostEvents(ctx, "ws-1") + if err != nil || len(costs) != 1 { + t.Fatalf("list costs: %v (len=%d)", err, len(costs)) + } + row := costs[0] + if row.CostSource == nil || *row.CostSource != models.CostSourceUnpriced { + t.Fatalf("CostSource = %v, want %q", row.CostSource, models.CostSourceUnpriced) + } + if row.Estimated { + t.Error("Estimated = true, want false: unpriced must not overwrite the adapter's own token-synthesis flag") + } +} + +// TestPromptUsage_TurnIDRecorded confirms turn_id threads all the way from +// the bus payload to the stored row when present. +func TestPromptUsage_TurnIDRecorded(t *testing.T) { + svc, eb := newTestServiceWithBus(t) + ctx := context.Background() + + createTestAgent(t, svc, "ws-1", "worker-turn") + insertTestTask(t, svc, "task-turn", "ws-1") + setTestTaskAssignee(t, svc, "task-turn", "worker-turn") + + event := bus.NewEvent(events.SessionPromptUsageUpdated, "test", map[string]interface{}{ + "task_id": "task-turn", + "session_id": "session-turn", + "model": "sonnet", + "turn_id": "turn-abc-123", + "usage_event_id": "usage-evt-abc-123", + "usage": map[string]interface{}{ + "input_tokens": 10, + "output_tokens": 5, + }, + }) + if err := eb.Publish(ctx, events.BuildSessionPromptUsageSubject("session-turn"), event); err != nil { + t.Fatalf("publish: %v", err) + } + + costs, err := svc.ListCostEvents(ctx, "ws-1") + if err != nil || len(costs) != 1 { + t.Fatalf("list costs: %v (len=%d)", err, len(costs)) + } + row := costs[0] + if row.TurnID == nil || *row.TurnID != "turn-abc-123" { + t.Errorf("TurnID = %v, want turn-abc-123", row.TurnID) + } + if row.UsageEventID == nil || *row.UsageEventID != "usage-evt-abc-123" { + t.Errorf("UsageEventID = %v, want usage-evt-abc-123", row.UsageEventID) + } +} + +// TestPromptUsage_DuplicateUsageEventIDIsIdempotent confirms redelivery of +// the same prompt-usage event (identical usage_event_id — e.g. an +// at-least-once event bus retry) does not double-record cost. +func TestPromptUsage_DuplicateUsageEventIDIsIdempotent(t *testing.T) { + svc, eb := newTestServiceWithBus(t) + ctx := context.Background() + + createTestAgent(t, svc, "ws-1", "worker-idem") + insertTestTask(t, svc, "task-idem", "ws-1") + setTestTaskAssignee(t, svc, "task-idem", "worker-idem") + + makeEvent := func() *bus.Event { + return bus.NewEvent(events.SessionPromptUsageUpdated, "test", map[string]interface{}{ + "task_id": "task-idem", + "session_id": "session-idem", + "model": "sonnet", + "usage_event_id": "usage-evt-dup-1", + "usage": map[string]interface{}{ + "input_tokens": 10, + "output_tokens": 5, + }, + }) + } + + if err := eb.Publish(ctx, events.BuildSessionPromptUsageSubject("session-idem"), makeEvent()); err != nil { + t.Fatalf("publish first: %v", err) + } + if err := eb.Publish(ctx, events.BuildSessionPromptUsageSubject("session-idem"), makeEvent()); err != nil { + t.Fatalf("publish redelivery: %v", err) + } + + costs, err := svc.ListCostEvents(ctx, "ws-1") + if err != nil { + t.Fatalf("list costs: %v", err) + } + if len(costs) != 1 { + t.Fatalf("cost count = %d, want 1 (redelivery must not double-record)", len(costs)) + } +} diff --git a/apps/backend/internal/office/service/event_subscribers_test.go b/apps/backend/internal/office/service/event_subscribers_test.go index 20f8d52683..9c86165436 100644 --- a/apps/backend/internal/office/service/event_subscribers_test.go +++ b/apps/backend/internal/office/service/event_subscribers_test.go @@ -518,12 +518,19 @@ func TestPromptUsage_RecordsCostEvent(t *testing.T) { t.Fatalf("cost event model = %q, want %q", costs[0].Model, "gpt-4o-mini") } // With no pricing lookup wired and no provider-reported cost, the row - // records 0 with estimated=true (Layer A miss + no Layer B). + // records 0 (Layer A miss + no Layer B), tagged cost_source=unpriced. + // Estimated tracks data.Usage.Estimated verbatim (unset here, so false) + // — it is a token-synthesis flag, distinct from cost_source, which is + // what actually carries the "we could not resolve a price" signal. See + // costContractVersion's v1->v2 doc comment in prompt_usage_cost.go. if costs[0].CostSubcents != 0 { t.Fatalf("cost_subcents = %d, want 0 (no pricing lookup wired)", costs[0].CostSubcents) } - if !costs[0].Estimated { - t.Fatal("expected estimated=true when both lookup layers miss") + if costs[0].Estimated { + t.Fatal("expected estimated=false: usage carried no estimated flag, and an unresolved price must not overwrite it") + } + if costs[0].CostSource == nil || *costs[0].CostSource != models.CostSourceUnpriced { + t.Fatalf("cost_source = %v, want %q", costs[0].CostSource, models.CostSourceUnpriced) } } diff --git a/apps/backend/internal/office/service/prompt_usage_cost.go b/apps/backend/internal/office/service/prompt_usage_cost.go new file mode 100644 index 0000000000..823383c92e --- /dev/null +++ b/apps/backend/internal/office/service/prompt_usage_cost.go @@ -0,0 +1,174 @@ +package service + +import ( + "context" + "time" + + "github.com/kandev/kandev/internal/office/costs" + "github.com/kandev/kandev/internal/office/models" + "github.com/kandev/kandev/internal/office/repository/sqlite" + "github.com/kandev/kandev/internal/office/shared" +) + +// costContractVersion is the in-band activation point for the cache-split / +// cost-provenance / turn-attribution columns (docs/specs/office/costs.md). +// The Rill cost extract has no schema versioning of its own, so a row +// written under a prior contract is distinguished by comparing +// cost_contract_version, not by a date an analyst has to be told out of +// band. Bump only if the contract's meaning changes again. +// +// v1 → v2: on the CostSourceUnpriced path, v1 forced Estimated=true +// regardless of data.Usage.Estimated, conflating "we could not resolve a +// price" with "the token counts themselves were synthesised" — two +// different signals this same contract introduced CostSource specifically +// to keep separate. v2 preserves data.Usage.Estimated verbatim on every +// path (see resolveCostForUsage); cost_source=unpriced alone now carries +// the pricing-failure signal. +const costContractVersion int64 = 2 + +// costResolution is resolveCostForUsage's output: the priced cost plus +// everything needed to record provenance on the row. Kept separate from +// models.CostEvent so this package's cost-resolution logic doesn't need to +// know about ID/session/task/occurred_at bookkeeping. +type costResolution struct { + costSubcents int64 + estimated bool + source models.CostSource + rates *shared.ModelPricing // nil unless source == CostSourceModelsDevList + catalogVersion string +} + +// resolveCostForUsage applies the Layer A / Layer B lookup and records +// which layer produced the dollar amount — CostSource, distinct from +// Estimated (a token-synthesis flag, not cost provenance). Layer A wins +// when the adapter forwarded a provider-reported cost sample, including an +// explicit zero +// (claude-acp's usage_update.cost.amount). Layer B (models.dev) is queried +// when a PricingLookup is wired; on miss or when no PricingLookup is +// configured the row is unpriced. Estimated is data.Usage.Estimated +// verbatim on every branch, including unpriced: whether the tokens were +// synthesised and whether a price could be resolved are independent facts, +// and cost_source=unpriced already carries the second one — see +// costContractVersion's v1→v2 doc comment. +func (s *Service) resolveCostForUsage(ctx context.Context, data PromptUsageData) costResolution { + if data.Usage.ProviderReportedCostPresent || data.Usage.ProviderReportedCostSubcents > 0 { + return costResolution{ + costSubcents: data.Usage.ProviderReportedCostSubcents, + estimated: data.Usage.Estimated, + source: models.CostSourceProviderReported, + } + } + if s.pricingLookup == nil || data.Model == "" { + return costResolution{estimated: data.Usage.Estimated, source: models.CostSourceUnpriced} + } + pricing, catalogVersion, ok := s.lookupPricingWithVersion(ctx, data.Model) + if !ok { + return costResolution{estimated: data.Usage.Estimated, source: models.CostSourceUnpriced} + } + cost := costs.CalculateCostSubcents( + data.Usage.InputTokens, + data.Usage.CachedReadTokens, + data.Usage.CachedWriteTokens, + data.Usage.OutputTokens, + costs.ModelPricing{ + InputPerMillion: pricing.InputPerMillion, + CachedReadPerMillion: pricing.CachedReadPerMillion, + CachedWritePerMillion: pricing.CachedWritePerMillion, + OutputPerMillion: pricing.OutputPerMillion, + }, + ) + return costResolution{ + costSubcents: cost, + estimated: data.Usage.Estimated, + source: models.CostSourceModelsDevList, + rates: &pricing, + catalogVersion: catalogVersion, + } +} + +// lookupPricingWithVersion resolves pricing and its catalogue version from +// one atomic snapshot when s.pricingLookup satisfies +// shared.PricingLookupWithVersion, so a concurrent background refresh can +// never pair one catalogue's rates with a different catalogue's version +// identifier on the stored row (docs/specs/office/costs.md). Falls back to two +// separate calls — accepting that narrower race — only for a PricingLookup +// implementation (e.g. a test double) that doesn't support the atomic form; +// CatalogVersion is optional there too, so a non-implementer simply reports +// no version. +func (s *Service) lookupPricingWithVersion( + ctx context.Context, model string, +) (shared.ModelPricing, string, bool) { + if withVersion, ok := s.pricingLookup.(shared.PricingLookupWithVersion); ok { + return withVersion.LookupForModelWithVersion(ctx, model) + } + pricing, ok := s.pricingLookup.LookupForModel(ctx, model) + if !ok { + return shared.ModelPricing{}, "", false + } + var version string + if versioner, ok := s.pricingLookup.(shared.PricingCatalogVersioner); ok { + version = versioner.CatalogVersion() + } + return pricing, version, true +} + +// buildCostEvent assembles the office_cost_events row for a prompt-usage +// update. +// +// The cache split (TokensCachedRead / TokensCachedWrite) is recorded only +// when data.Usage.Estimated is false. codex-acp's ACP bridge (and any +// future adapter with no per-turn usage frame) synthesises InputTokens from +// context-occupancy growth with Estimated=true and never sets the cache +// fields — a NULL split there is honest; 0 would silently claim "no cache +// activity" for an agent that never reported one. TokensCachedIn keeps its +// original read+write sum on every row regardless, so existing consumers +// (the tree-holds rollup, card 2faa29da's task_sessions fix) are +// unaffected. +func buildCostEvent( + data PromptUsageData, fields *sqlite.TaskExecutionFields, projectID, provider string, + resolution costResolution, +) *models.CostEvent { + event := &models.CostEvent{ + SessionID: data.SessionID, + TaskID: data.TaskID, + AgentProfileID: fields.AssigneeAgentProfileID, + ProjectID: projectID, + Model: data.Model, + Provider: provider, + TokensIn: data.Usage.InputTokens, + TokensCachedIn: data.Usage.CachedReadTokens + data.Usage.CachedWriteTokens, + TokensOut: data.Usage.OutputTokens, + CostSubcents: resolution.costSubcents, + Estimated: resolution.estimated, + CostSource: &resolution.source, + OccurredAt: time.Now().UTC(), + } + contractVersion := costContractVersion + event.CostContractVersion = &contractVersion + + if !data.Usage.Estimated { + cachedRead := data.Usage.CachedReadTokens + cachedWrite := data.Usage.CachedWriteTokens + event.TokensCachedRead = &cachedRead + event.TokensCachedWrite = &cachedWrite + } + if resolution.rates != nil { + event.RateInputPerMillion = &resolution.rates.InputPerMillion + event.RateCachedReadPerMillion = &resolution.rates.CachedReadPerMillion + event.RateCachedWritePerMillion = &resolution.rates.CachedWritePerMillion + event.RateOutputPerMillion = &resolution.rates.OutputPerMillion + } + if resolution.catalogVersion != "" { + version := resolution.catalogVersion + event.PricingCatalogVersion = &version + } + if data.TurnID != "" { + turnID := data.TurnID + event.TurnID = &turnID + } + if data.UsageEventID != "" { + usageEventID := data.UsageEventID + event.UsageEventID = &usageEventID + } + return event +} diff --git a/apps/backend/internal/office/shared/interfaces.go b/apps/backend/internal/office/shared/interfaces.go index 6081c72b88..7936e5a373 100644 --- a/apps/backend/internal/office/shared/interfaces.go +++ b/apps/backend/internal/office/shared/interfaces.go @@ -5,6 +5,8 @@ import ( "errors" "time" + "github.com/jmoiron/sqlx" + "github.com/kandev/kandev/internal/office/models" "github.com/kandev/kandev/internal/workflow/engine" ) @@ -150,9 +152,61 @@ type PricingLookup interface { LookupForModel(ctx context.Context, modelID string) (ModelPricing, bool) } +// PricingCatalogVersioner is an optional capability a PricingLookup +// implementation may satisfy to report an "as-of" identifier for the +// pricing data it served — recorded on CostEvent.PricingCatalogVersion so a +// models.dev-list-priced row can be traced back to the catalogue state that +// produced it. Deliberately separate from PricingLookup (rather than +// widening it) so existing implementers and test fakes are unaffected; +// callers type-assert and treat a missing implementation as "no version +// available" (NULL column), not an error. +type PricingCatalogVersioner interface { + // CatalogVersion returns an identifier for the currently-served + // pricing data, or "" when none is available yet (e.g. cold cache, + // nothing loaded). models.dev's dataset carries no version field of + // its own, so implementations report the load/fetch time instead. + CatalogVersion() string +} + +// PricingLookupWithVersion is an optional capability a PricingLookup +// implementation may satisfy to return pricing and its catalogue version +// from one atomic snapshot. Calling LookupForModel and CatalogVersion +// separately takes two independent lock acquisitions; a background refresh +// can install a new catalogue in between and pair one catalogue's rates +// with a different catalogue's version identifier on the stored row — a +// provenance column that lies, which is the exact failure class +// CostEvent.CostSource exists to eliminate. Callers type-assert and prefer +// this over the separate calls whenever both values are needed together. +type PricingLookupWithVersion interface { + // LookupForModelWithVersion behaves like PricingLookup.LookupForModel + // but also returns the catalogue version that produced the pricing, + // read from the same snapshot so the two can never describe different + // catalogue states. + LookupForModelWithVersion(ctx context.Context, modelID string) (pricing ModelPricing, version string, ok bool) +} + // SessionUsageWriter increments the cumulative tokens/cost columns on // task_sessions when a cost event lands. Implemented by the task repo. type SessionUsageWriter interface { IncrementTaskSessionUsage(ctx context.Context, sessionID string, tokensIn, tokensCachedIn, tokensOut, costSubcents int64) error } + +// SessionUsageWriterTx is an optional capability a SessionUsageWriter +// implementation may satisfy to run the increment inside a caller-supplied +// transaction, so it can be made atomic with the office_cost_events insert +// that produced the deltas (docs/specs/office/costs.md, PR #2606 review): +// without this, a rollup-increment failure after a successful cost-event +// insert is logged and swallowed, and any later redelivery of the same +// completion is caught by the usage_event_id unique index and dropped as a +// duplicate before the rollup gets another chance — a permanent desync with +// no recovery path. Callers type-assert and fall back to the plain +// (non-atomic) SessionUsageWriter method when unavailable — e.g. a +// simplified test double that doesn't need transactional coupling. Every +// production wiring (backendapp's SetSessionUsageWriter(repos.Task)) +// satisfies it, since office_cost_events and task_sessions share one +// underlying *sqlx.DB (see internal/backendapp/storage.go). +type SessionUsageWriterTx interface { + IncrementTaskSessionUsageTx(ctx context.Context, tx *sqlx.Tx, sessionID string, + tokensIn, tokensCachedIn, tokensOut, costSubcents int64) error +} diff --git a/apps/backend/internal/orchestrator/event_handlers_agent.go b/apps/backend/internal/orchestrator/event_handlers_agent.go index 0e21402bb6..01d0fcb47a 100644 --- a/apps/backend/internal/orchestrator/event_handlers_agent.go +++ b/apps/backend/internal/orchestrator/event_handlers_agent.go @@ -543,6 +543,13 @@ func (s *Service) handleAgentReady(ctx context.Context, data watcher.AgentEventD // later, unrelated provider overload starts its backoff fresh at attempt 1. s.resetTransientRetry(data.SessionID) + // Snapshot the turn this event is about to close for the office cost + // subscriber's benefit: publishPromptUsage's complete-stream frame for + // this same completion is published (and processed) after this handler + // returns, by which point completeTurnForSession below has already + // removed it from activeTurns. See markReadyTurn's doc comment. + s.markReadyTurn(data.SessionID, data.AgentExecutionID, data.PromptGeneration, turnAtEventFire) + // Complete the current turn s.completeTurnForSession(ctx, data.SessionID) diff --git a/apps/backend/internal/orchestrator/event_handlers_streaming.go b/apps/backend/internal/orchestrator/event_handlers_streaming.go index fdb2d306ef..c943ea2e1b 100644 --- a/apps/backend/internal/orchestrator/event_handlers_streaming.go +++ b/apps/backend/internal/orchestrator/event_handlers_streaming.go @@ -6,6 +6,7 @@ import ( "strings" "time" + "github.com/google/uuid" "go.uber.org/zap" "github.com/kandev/kandev/internal/agent/runtime/lifecycle" @@ -21,6 +22,12 @@ import ( const sessionModelConfigKey = "model" +// usageEventIDNamespace seeds the deterministic UUID usageEventIDFor derives +// for a prompt-usage completion. Arbitrary but fixed — any stable value +// works since it only needs to be consistent across process restarts, never +// shared with another namespace. +var usageEventIDNamespace = uuid.MustParse("2f6a6f8c-6c1b-4b8a-9e3e-7a6d2c5b9f10") + // handleAgentStreamEvent handles agent stream events (tool calls, message chunks, etc.) func (s *Service) handleAgentStreamEvent(ctx context.Context, payload *lifecycle.AgentStreamEventPayload) { if payload == nil || payload.Data == nil { @@ -228,10 +235,43 @@ func (s *Service) streamEventIsStalePrompt(payload *lifecycle.AgentStreamEventPa ) } -func (s *Service) completeTurnForStreamEvent(ctx context.Context, payload *lifecycle.AgentStreamEventPayload) { +func (s *Service) completeTurnForStreamEvent( + ctx context.Context, + payload *lifecycle.AgentStreamEventPayload, + capturedTurnIDs ...string, +) { if payload == nil { return } + var capturedTurnID string + if len(capturedTurnIDs) > 0 { + capturedTurnID = capturedTurnIDs[0] + } + if capturedTurnID != "" { + // A durable turn ID is authoritative. If the event is stale and an + // accepted successor exists, preserve that successor while settling + // the predecessor; otherwise close only the captured turn so a late + // completion cannot sweep an unrelated active turn. + stale := s.streamEventIsStalePrompt(payload) + successorTurnID := s.acceptedDispatchSuccessorTurn(payload.SessionID) + if successorTurnID != "" && successorTurnID != capturedTurnID { + stale = true + } + if stale && s.acceptedDispatchInFlight(payload.SessionID) { + s.completeTurnForTaskSessionWithSuccessorPolicy(ctx, payload.TaskID, payload.SessionID, true) + return + } + if err := s.completeTurnForTaskSessionCheckedOwned(ctx, payload.TaskID, payload.SessionID, capturedTurnID); err != nil { + s.logger.Warn("failed to complete stream event's captured turn", + zap.String("session_id", payload.SessionID), + zap.String("turn_id", capturedTurnID), + zap.Error(err)) + } + if successorTurnID == capturedTurnID { + s.clearAcceptedQueuedDispatch(payload.SessionID) + } + return + } s.completeTurnForTaskSessionWithSuccessorPolicy( ctx, payload.TaskID, @@ -1455,6 +1495,164 @@ func (s *Service) terminalExecutionMarker(sessionID, executionID string) (termin return marker, true } +type readyTurnMark struct { + turnID string + expiresAt time.Time +} + +func readyTurnKey(sessionID, executionID string, promptGeneration uint64) string { + return fmt.Sprintf("%s\x00%s\x00%d", sessionID, executionID, promptGeneration) +} + +// readyTurnZeroGenKey is deliberately generation-less: promptGeneration==0 +// means the transport carries no generation tracking at all (see +// finishPromptCompletion / claimPromptCompletion's early return), so every +// completion on this (session, execution) shares the same identity and must +// be threaded through the FIFO queue in s.readyTurnMarksZeroGen instead of +// the single-slot s.readyTurnMarks map. +func readyTurnZeroGenKey(sessionID, executionID string) string { + return fmt.Sprintf("%s\x00%s", sessionID, executionID) +} + +// markReadyTurn records the turn ID handleAgentReady confirmed for +// (sessionID, executionID, promptGeneration), just before completeTurnForSession +// closes it and removes it from activeTurns. handleCompleteEventMarkState +// (lifecycle package) publishes agent.ready and closes the turn before the +// complete-stream frame for the same completion is published, on EVERY +// transport, including promptGeneration==0 (generation-less) completions — +// so this mark is required there too, not skippable. Because generation-less +// completions share one key per (session, execution) with no generation to +// disambiguate them, they queue FIFO in readyTurnMarksZeroGen instead of the +// single-slot readyTurnMarks map: ready and complete-stream for the same +// completion are published strictly in that order per execution, so pending +// marks and pending completions correlate 1:1 in arrival order. +// +// This ordering is airtight on the default in-memory event bus, where +// Publish delivers to a subject's subscriber synchronously before returning +// (see internal/events/bus/memory.go). It is best-effort, not guaranteed, on +// a NATS-backed bus (opt-in via cfg.NATS.URL): AgentReady and the +// agent.stream.* wildcard are different subjects, and this in-memory map is +// per-process. The durable TurnID on the completion payload closes that +// cross-subject/cross-instance gap; this mark remains a compatibility fallback +// for older producers and completion paths without a captured ID. +func (s *Service) markReadyTurn(sessionID, executionID string, promptGeneration uint64, turnID string) { + if sessionID == "" || executionID == "" || turnID == "" { + return + } + expiresAt := time.Now().Add(completedExecutionRetention) + if promptGeneration == 0 { + s.pushZeroGenReadyTurnMark(sessionID, executionID, readyTurnMark{turnID: turnID, expiresAt: expiresAt}) + return + } + key := readyTurnKey(sessionID, executionID, promptGeneration) + s.readyTurnMarks.Store(key, readyTurnMark{turnID: turnID, expiresAt: expiresAt}) + time.AfterFunc(completedExecutionRetention, func() { + s.deleteReadyTurnMarkIfExpired(key, expiresAt) + }) +} + +// takeReadyTurnMark consumes (and removes) the turn ID markReadyTurn recorded +// for this completion, if any. A miss is expected whenever the completion +// never went through handleAgentReady's synchronous ready path (or the mark +// already expired); the caller falls back to a live lookup or a terminal- +// execution snapshot in that case. +func (s *Service) takeReadyTurnMark(sessionID, executionID string, promptGeneration uint64) (string, bool) { + if sessionID == "" || executionID == "" { + return "", false + } + if promptGeneration == 0 { + return s.popZeroGenReadyTurnMark(sessionID, executionID) + } + key := readyTurnKey(sessionID, executionID, promptGeneration) + value, ok := s.readyTurnMarks.LoadAndDelete(key) + if !ok { + return "", false + } + mark, ok := value.(readyTurnMark) + if !ok || time.Now().After(mark.expiresAt) { + return "", false + } + return mark.turnID, true +} + +func (s *Service) deleteReadyTurnMarkIfExpired(key string, expiresAt time.Time) { + value, ok := s.readyTurnMarks.Load(key) + if !ok { + return + } + current, ok := value.(readyTurnMark) + if !ok || !current.expiresAt.After(expiresAt) { + s.readyTurnMarks.Delete(key) + } +} + +// pushZeroGenReadyTurnMark appends a generation-less ready-turn mark to the +// FIFO queue for (sessionID, executionID). Guarded by +// readyTurnMarksZeroGenMu: sync.Map has no atomic append, and multiple +// pending marks for the same key are the expected case here (unlike the +// generation-keyed map, where each key holds at most one). +func (s *Service) pushZeroGenReadyTurnMark(sessionID, executionID string, mark readyTurnMark) { + key := readyTurnZeroGenKey(sessionID, executionID) + s.readyTurnMarksZeroGenMu.Lock() + if s.readyTurnMarksZeroGen == nil { + s.readyTurnMarksZeroGen = make(map[string][]readyTurnMark) + } + s.readyTurnMarksZeroGen[key] = append(s.readyTurnMarksZeroGen[key], mark) + s.readyTurnMarksZeroGenMu.Unlock() + time.AfterFunc(completedExecutionRetention, func() { + s.pruneExpiredZeroGenReadyTurnMarks(key) + }) +} + +// popZeroGenReadyTurnMark consumes the oldest live mark queued for +// (sessionID, executionID), discarding any expired entries ahead of it. +func (s *Service) popZeroGenReadyTurnMark(sessionID, executionID string) (string, bool) { + key := readyTurnZeroGenKey(sessionID, executionID) + s.readyTurnMarksZeroGenMu.Lock() + defer s.readyTurnMarksZeroGenMu.Unlock() + queue := s.readyTurnMarksZeroGen[key] + now := time.Now() + for len(queue) > 0 { + mark := queue[0] + queue = queue[1:] + if now.After(mark.expiresAt) { + continue + } + if len(queue) == 0 { + delete(s.readyTurnMarksZeroGen, key) + } else { + s.readyTurnMarksZeroGen[key] = queue + } + return mark.turnID, true + } + delete(s.readyTurnMarksZeroGen, key) + return "", false +} + +// pruneExpiredZeroGenReadyTurnMarks drops expired entries from the front of +// (sessionID+executionID)'s queue so an unconsumed mark cannot grow the map +// unbounded, mirroring completedExecutions' expiry pattern. +func (s *Service) pruneExpiredZeroGenReadyTurnMarks(key string) { + s.readyTurnMarksZeroGenMu.Lock() + defer s.readyTurnMarksZeroGenMu.Unlock() + queue := s.readyTurnMarksZeroGen[key] + if len(queue) == 0 { + return + } + now := time.Now() + kept := queue[:0] + for _, mark := range queue { + if now.Before(mark.expiresAt) { + kept = append(kept, mark) + } + } + if len(kept) == 0 { + delete(s.readyTurnMarksZeroGen, key) + return + } + s.readyTurnMarksZeroGen[key] = kept +} + func (s *Service) currentTurnIDForSession(ctx context.Context, sessionID string) string { if sessionID == "" { return "" @@ -1994,35 +2192,66 @@ func (s *Service) handleCompleteStreamEvent(ctx context.Context, payload *lifecy s.storeResumeToken(ctx, payload.TaskID, payload.SessionID, payload.ExecutionID, payload.Data.ACPSessionID, lastMsgUUID) } - s.publishPromptUsage(ctx, payload, session) + // The lifecycle completion payload carries the durable turn captured before + // AgentReady can admit a successor prompt. Older producers do not include + // it, so retain the generation-keyed (or generation-less FIFO) ready mark as + // a compatibility fallback. A miss also means this completion never went + // through handleAgentReady's ready path (e.g. an error/interrupt that + // completed the execution directly), so fall back to whichever snapshot the + // branch has: terminalMarker.turnID (captured by markTerminalExecution at + // agent.completed) for terminal completions, or a live active-turn lookup + // for non-terminal ones. + completionTurnID := payload.Data.TurnID + if completionTurnID == "" { + var ok bool + completionTurnID, ok = s.takeReadyTurnMark(payload.SessionID, payload.ExecutionID, payload.Data.PromptGeneration) + if !ok { + if terminalCompleteStream { + completionTurnID = terminalMarker.turnID + } else { + completionTurnID = s.currentTurnIDForSession(ctx, payload.SessionID) + } + } + } + s.publishPromptUsage(ctx, payload, session, completionTurnID) if terminalCompleteStream { - s.saveAgentTextForTurn(ctx, payload, terminalMarker.turnID) - s.publishAgentPlanForTurn(ctx, payload, terminalMarker.turnID, false) - s.persistTurnPromptMetadataForTurn(ctx, payload, session, terminalMarker.turnID) - if terminalMarker.turnID != "" { - s.publishAgentTurnCompleteForTurn(ctx, payload, terminalMarker.turnID) + terminalTurnID := terminalMarker.turnID + if payload.Data.TurnID != "" { + terminalTurnID = payload.Data.TurnID + } + s.saveAgentTextForTurn(ctx, payload, terminalTurnID) + s.publishAgentPlanForTurn(ctx, payload, terminalTurnID, false) + s.persistTurnPromptMetadataForTurn(ctx, payload, session, terminalTurnID) + if terminalTurnID != "" { + s.publishAgentTurnCompleteForTurn(ctx, payload, terminalTurnID) } s.detachClarificationWaiters(ctx, payload.SessionID) s.logger.Debug("complete stream from terminal execution flushed final data; skipping active turn and runtime reconciliation", zap.String("task_id", payload.TaskID), zap.String("session_id", payload.SessionID), zap.String("agent_execution_id", payload.ExecutionID), - zap.String("turn_id", terminalMarker.turnID)) + zap.String("turn_id", terminalTurnID)) return } - s.saveAgentTextIfPresent(ctx, payload) - s.publishAgentPlanIfPresent(ctx, payload) - s.persistTurnPromptMetadata(ctx, payload, session) - s.completeTurnForStreamEvent(ctx, payload) + if completionTurnID != "" { + s.saveAgentTextForTurn(ctx, payload, completionTurnID) + s.publishAgentPlanForTurn(ctx, payload, completionTurnID, false) + s.persistTurnPromptMetadataForTurn(ctx, payload, session, completionTurnID) + } else { + s.saveAgentTextIfPresent(ctx, payload) + s.publishAgentPlanIfPresent(ctx, payload) + s.persistTurnPromptMetadata(ctx, payload, session) + } + s.completeTurnForStreamEvent(ctx, payload, completionTurnID) // Publish agent turn message event so the office comment bridge can // auto-post the agent's response as a task comment. Published here // (not in saveAgentTextIfPresent) because for streaming agents the // text is drained by message_chunk events and Data.Text is empty at // complete time. - s.publishAgentTurnComplete(ctx, payload) + s.publishAgentTurnCompleteForTurn(ctx, payload, completionTurnID) // Detach any pending clarifications so WaitForResponse unblocks while the // overlay stays interactive for a deferred answer via the event fallback path. @@ -2260,15 +2489,49 @@ func (s *Service) publishAgentPlanForTurn(ctx context.Context, payload *lifecycl } } -// publishPromptUsage broadcasts prompt token usage to the WebSocket for the frontend. -// Model and agent type (CLI engine slug) come from payload first; when absent -// (which is the common case — CurrentModelID only travels on session_models -// frames) we fall back to the session's AgentProfileSnapshot, populated at -// session creation and refreshed by persistSessionModel on ACP model updates. +// usageEventIDFor derives the office cost subscriber's idempotency key from +// immutable upstream identity — session, execution, and prompt generation — +// instead of minting a fresh random value on every call. A minted-per-call +// key can only dedup literal redelivery of the identical *bus.Event object, +// which neither event bus provides (see internal/events/bus/{memory,nats}.go +// — the memory bus delivers once synchronously with no retry, and the NATS +// bus is plain core pub/sub, not JetStream). The real duplicate source is +// the SAME underlying completion frame reaching publishPromptUsage twice — +// e.g. a reconnecting WS client replaying a buffered stream event. +// +// promptGeneration==0 means this completion carries no generation tracking +// at all (see claimPromptCompletion's early return in the lifecycle +// package); deriving a key from (session, execution, 0) there would collide +// across genuinely distinct turns on a generation-less transport and +// silently under-count cost, which is worse than the duplicate-row bug this +// fixes. Fall back to a random id in that narrow case — unchanged from +// prior behavior there. +func usageEventIDFor(sessionID, executionID string, promptGeneration uint64) string { + if sessionID == "" || executionID == "" || promptGeneration == 0 { + return uuid.New().String() + } + name := fmt.Sprintf("%s\x00%s\x00%d", sessionID, executionID, promptGeneration) + return uuid.NewSHA1(usageEventIDNamespace, []byte(name)).String() +} + +// publishPromptUsage broadcasts prompt token usage to the WebSocket for the +// frontend and to the office cost subscriber. Model and agent type (CLI +// engine slug) come from payload first; when absent (which is the common +// case — CurrentModelID only travels on session_models frames) we fall back +// to the session's AgentProfileSnapshot, populated at session creation and +// refreshed by persistSessionModel on ACP model updates. +// +// turnID is resolved by the caller (handleCompleteStreamEvent), not here: +// the terminal-execution snapshot and the live active-turn lookup are both +// call-site concerns. usageEventID is derived here, once, at the single +// publish site by usageEventIDFor — that is what makes it a stable +// idempotency key across a republished frame; a downstream consumer +// deriving its own would defeat the point. func (s *Service) publishPromptUsage( ctx context.Context, payload *lifecycle.AgentStreamEventPayload, session *models.TaskSession, + turnID string, ) { sessionID := payload.SessionID if sessionID == "" || s.eventBus == nil || payload.Data.Usage == nil { @@ -2285,6 +2548,10 @@ func (s *Service) publishPromptUsage( Model: model, Usage: payload.Data.Usage, Timestamp: time.Now().UTC().Format(time.RFC3339), + TurnID: turnID, + UsageEventID: usageEventIDFor( + sessionID, payload.ExecutionID, payload.Data.PromptGeneration, + ), } subject := events.BuildSessionPromptUsageSubject(sessionID) _ = s.eventBus.Publish(ctx, subject, bus.NewEvent(events.SessionPromptUsageUpdated, "orchestrator", eventPayload)) @@ -2398,6 +2665,7 @@ func promptUsageMetadata(usage *streams.PromptUsage) map[string]interface{} { "thought_tokens": usage.ThoughtTokens, "total_tokens": usage.TotalTokens, "provider_reported_cost_subcents": usage.ProviderReportedCostSubcents, + "provider_reported_cost_present": usage.ProviderReportedCostPresent, "estimated": usage.Estimated, } } diff --git a/apps/backend/internal/orchestrator/event_handlers_streaming_test.go b/apps/backend/internal/orchestrator/event_handlers_streaming_test.go index 133fd14c4a..18c1e52eff 100644 --- a/apps/backend/internal/orchestrator/event_handlers_streaming_test.go +++ b/apps/backend/internal/orchestrator/event_handlers_streaming_test.go @@ -1293,6 +1293,404 @@ func TestCompleteStreamFromCompletedExecutionPublishesTerminalTurn(t *testing.T) "late terminal complete publish must identify the completed turn") } +func findPromptUsageEvent(t *testing.T, eb *recordingEventBus) *lifecycle.SessionPromptUsageEventPayload { + t.Helper() + for _, rec := range eb.events { + if rec.subject != events.BuildSessionPromptUsageSubject("s1") { + continue + } + payload, ok := rec.event.Data.(lifecycle.SessionPromptUsageEventPayload) + require.True(t, ok, "session_prompt_usage.updated event carried an unexpected payload type") + return &payload + } + return nil +} + +// TestPublishPromptUsage_NonTerminalCompletionUsesReadyTurnSnapshot is the +// regression test for the ordinary (non-terminal) completion path: the agent +// stays running, so agent.ready — not agent.completed — closes the turn. +// handleAgentReady publishes agent.ready synchronously and closes the turn via +// completeTurnForSession before the complete-stream frame for the same +// completion is ever processed (see markReadyTurn's and +// handleCompleteStreamEvent's doc comments), so a live active-turn lookup at +// that point finds nothing and would store turn_id NULL. +func TestPublishPromptUsage_NonTerminalCompletionUsesReadyTurnSnapshot(t *testing.T) { + ctx := context.Background() + repo := setupTestRepo(t) + seedSession(t, repo, "t1", "s1", "step1") + + stepGetter := newMockStepGetter() + stepGetter.steps["step1"] = &wfmodels.WorkflowStep{ + ID: "step1", WorkflowID: "wf1", Name: "Step 1", Position: 0, + } + taskRepo := newMockTaskRepo() + seedMockTaskState(taskRepo, "t1", v1.TaskStateInProgress) + agentMgr := &mockAgentManager{isAgentRunning: true} + agentMgr.currentPromptExecutionID = "exec-1" + agentMgr.currentPromptGeneration.Store(7) + svc := createTestServiceWithAgent(repo, stepGetter, taskRepo, agentMgr) + svc.turnService = &repoTurnService{repo: repo} + eb := &recordingEventBus{} + svc.eventBus = eb + + turn, err := svc.turnService.StartTurn(ctx, "s1") + require.NoError(t, err) + + // The ready event closes the turn synchronously — before this test (or + // production) ever gets to publish the complete-stream frame below. + svc.handleAgentReady(ctx, watcher.AgentEventData{ + TaskID: "t1", SessionID: "s1", AgentExecutionID: "exec-1", PromptGeneration: 7, + }) + + active, err := svc.turnService.GetActiveTurn(ctx, "s1") + require.NoError(t, err) + require.Nil(t, active, "handleAgentReady must have already closed the turn") + + svc.handleAgentStreamEvent(ctx, &lifecycle.AgentStreamEventPayload{ + TaskID: "t1", + SessionID: "s1", + ExecutionID: "exec-1", + Data: &lifecycle.AgentStreamEventData{ + Type: agentEventComplete, + PromptGeneration: 7, + Usage: &streams.PromptUsage{InputTokens: 10, OutputTokens: 5}, + }, + }) + + usageEvent := findPromptUsageEvent(t, eb) + require.NotNil(t, usageEvent, "expected a session_prompt_usage.updated event to be published") + require.Equal(t, turn.ID, usageEvent.TurnID, + "non-terminal completion must carry the turn id the ready event just closed, not NULL") +} + +// TestPublishPromptUsage_CompletionPayloadTurnIDOwnsTurn verifies that a +// completion's durable turn identity wins over the currently active turn. +// This is the ordering that occurs when a queued successor starts before the +// completion frame crosses a NATS subject or instance boundary. +func TestPublishPromptUsage_CompletionPayloadTurnIDOwnsTurn(t *testing.T) { + ctx := context.Background() + repo := setupTestRepo(t) + seedSession(t, repo, "t1", "s1", "step1") + + stepGetter := newMockStepGetter() + stepGetter.steps["step1"] = &wfmodels.WorkflowStep{ + ID: "step1", WorkflowID: "wf1", Name: "Step 1", Position: 0, + } + taskRepo := newMockTaskRepo() + seedMockTaskState(taskRepo, "t1", v1.TaskStateInProgress) + svc := createTestService(repo, stepGetter, taskRepo) + svc.turnService = &repoTurnService{repo: repo} + eb := &recordingEventBus{} + svc.eventBus = eb + + completedTurn, err := svc.turnService.StartTurn(ctx, "s1") + require.NoError(t, err) + svc.completeTurnForSession(ctx, "s1") + successorTurn, err := svc.turnService.StartTurn(ctx, "s1") + require.NoError(t, err) + + svc.handleAgentStreamEvent(ctx, &lifecycle.AgentStreamEventPayload{ + TaskID: "t1", + SessionID: "s1", + ExecutionID: "exec-1", + Data: &lifecycle.AgentStreamEventData{ + Type: agentEventComplete, + TurnID: completedTurn.ID, + Usage: &streams.PromptUsage{InputTokens: 10, OutputTokens: 5}, + }, + }) + + usageEvent := findPromptUsageEvent(t, eb) + require.NotNil(t, usageEvent, "expected a session_prompt_usage.updated event to be published") + require.Equal(t, completedTurn.ID, usageEvent.TurnID, + "completion payload must retain its captured turn id") + + active, err := svc.turnService.GetActiveTurn(ctx, "s1") + require.NoError(t, err) + require.Equal(t, successorTurn.ID, active.ID, + "a late completion must not close a successor turn") +} + +// TestMarkReadyTurn_ZeroGenerationQueuesFIFO is the unit-level regression +// test for R2-F3: markReadyTurn's old promptGeneration==0 early return +// rested on the false premise that generation-less completions never +// publish agent.ready synchronously (see markReadyTurn's doc comment for the +// actual mechanism — handleCompleteEventMarkState calls MarkReady +// unconditionally). Two pending marks on the same (session, execution) with +// no generation to disambiguate them must not collide — they queue FIFO and +// come back out in the order they were recorded. +func TestMarkReadyTurn_ZeroGenerationQueuesFIFO(t *testing.T) { + svc := &Service{} + + svc.markReadyTurn("s1", "exec-1", 0, "turn-A") + svc.markReadyTurn("s1", "exec-1", 0, "turn-B") + + first, ok := svc.takeReadyTurnMark("s1", "exec-1", 0) + require.True(t, ok, "expected a mark for the first pending completion") + require.Equal(t, "turn-A", first, "FIFO: the first mark recorded must be the first consumed") + + second, ok := svc.takeReadyTurnMark("s1", "exec-1", 0) + require.True(t, ok, "expected a mark for the second pending completion") + require.Equal(t, "turn-B", second) + + _, ok = svc.takeReadyTurnMark("s1", "exec-1", 0) + require.False(t, ok, "queue must be drained after both marks are consumed") + + // A different execution on the same session must not share the queue. + svc.markReadyTurn("s1", "exec-2", 0, "turn-C") + _, ok = svc.takeReadyTurnMark("s1", "exec-1", 0) + require.False(t, ok, "exec-1's (already-drained) queue must not see exec-2's mark") + turnC, ok := svc.takeReadyTurnMark("s1", "exec-2", 0) + require.True(t, ok) + require.Equal(t, "turn-C", turnC) +} + +// TestPublishPromptUsage_NonTerminalCompletionZeroGenerationUsesReadyTurnSnapshot +// is the integration-level regression test for R2-F3: the generation-0 +// (generation-less transport) sibling of +// TestPublishPromptUsage_NonTerminalCompletionUsesReadyTurnSnapshot. Before +// the fix, markReadyTurn no-op'd for promptGeneration==0, so this scenario +// always stored turn_id NULL even though agent.ready closes the turn here +// exactly as it does for a generation-tracked completion. +func TestPublishPromptUsage_NonTerminalCompletionZeroGenerationUsesReadyTurnSnapshot(t *testing.T) { + ctx := context.Background() + repo := setupTestRepo(t) + seedSession(t, repo, "t1", "s1", "step1") + + stepGetter := newMockStepGetter() + stepGetter.steps["step1"] = &wfmodels.WorkflowStep{ + ID: "step1", WorkflowID: "wf1", Name: "Step 1", Position: 0, + } + taskRepo := newMockTaskRepo() + seedMockTaskState(taskRepo, "t1", v1.TaskStateInProgress) + agentMgr := &mockAgentManager{isAgentRunning: true} + agentMgr.currentPromptExecutionID = "exec-1" + svc := createTestServiceWithAgent(repo, stepGetter, taskRepo, agentMgr) + svc.turnService = &repoTurnService{repo: repo} + eb := &recordingEventBus{} + svc.eventBus = eb + + turn, err := svc.turnService.StartTurn(ctx, "s1") + require.NoError(t, err) + + svc.handleAgentReady(ctx, watcher.AgentEventData{ + TaskID: "t1", SessionID: "s1", AgentExecutionID: "exec-1", PromptGeneration: 0, + }) + + active, err := svc.turnService.GetActiveTurn(ctx, "s1") + require.NoError(t, err) + require.Nil(t, active, "handleAgentReady must have already closed the turn") + + svc.handleAgentStreamEvent(ctx, &lifecycle.AgentStreamEventPayload{ + TaskID: "t1", + SessionID: "s1", + ExecutionID: "exec-1", + Data: &lifecycle.AgentStreamEventData{ + Type: agentEventComplete, + Usage: &streams.PromptUsage{InputTokens: 10, OutputTokens: 5}, + }, + }) + + usageEvent := findPromptUsageEvent(t, eb) + require.NotNil(t, usageEvent, "expected a session_prompt_usage.updated event to be published") + require.Equal(t, turn.ID, usageEvent.TurnID, + "generation-less non-terminal completion must carry the turn id the ready event just closed, not NULL") +} + +// TestPublishPromptUsage_TerminalCompletionWithoutReadyMarkFallsBackToTerminalMarker +// covers the terminal (agent.completed) fallback path for a completion that +// never went through handleAgentReady's synchronous ready path at all — e.g. +// a crash mid-prompt, where finishPromptCompletion never runs and no ready +// mark is ever recorded. There, markTerminalExecution's live-turn snapshot +// (captured before completeTurnForSession closes it) is the only source of +// truth, so this test deliberately keeps that ordering. +func TestPublishPromptUsage_TerminalCompletionWithoutReadyMarkFallsBackToTerminalMarker(t *testing.T) { + ctx := context.Background() + repo := setupTestRepo(t) + seedSession(t, repo, "t1", "s1", "") + + svc := createTestService(repo, newMockStepGetter(), newMockTaskRepo()) + svc.turnService = &repoTurnService{repo: repo} + eb := &recordingEventBus{} + svc.eventBus = eb + + turn, err := svc.turnService.StartTurn(ctx, "s1") + require.NoError(t, err) + svc.markExecutionCompleted("s1", "exec-1") + svc.completeTurnForSession(ctx, "s1") + + svc.handleAgentStreamEvent(ctx, &lifecycle.AgentStreamEventPayload{ + TaskID: "t1", + SessionID: "s1", + ExecutionID: "exec-1", + Data: &lifecycle.AgentStreamEventData{ + Type: agentEventComplete, + Usage: &streams.PromptUsage{InputTokens: 10, OutputTokens: 5}, + }, + }) + + usageEvent := findPromptUsageEvent(t, eb) + require.NotNil(t, usageEvent, "expected a session_prompt_usage.updated event to be published") + require.Equal(t, turn.ID, usageEvent.TurnID, + "terminal completion with no ready mark must fall back to the terminal marker's snapshotted turn id") +} + +// TestPublishPromptUsage_TerminalCompletionAfterReadyUsesReadyTurnSnapshot is +// the R2-F1 regression test: the REALISTIC terminal-completion ordering, +// where agent.ready fires (and closes the turn) before the process later +// exits and agent.completed marks the execution terminal. +// finishPromptCompletion (lifecycle package) publishes agent.ready +// synchronously on EVERY successful prompt completion, independent of +// whether the process subsequently exits — so by the time +// markExecutionCompleted runs, the turn markTerminalExecution would snapshot +// via a live lookup is already closed (returns ""), and the correct turn id +// is only available from the ready mark handleAgentReady recorded. Before +// the R2-F1 fix, the terminal branch of handleCompleteStreamEvent read +// terminalMarker.turnID directly and never consulted that mark, so this +// scenario stored turn_id NULL. +func TestPublishPromptUsage_TerminalCompletionAfterReadyUsesReadyTurnSnapshot(t *testing.T) { + ctx := context.Background() + repo := setupTestRepo(t) + seedSession(t, repo, "t1", "s1", "step1") + + stepGetter := newMockStepGetter() + stepGetter.steps["step1"] = &wfmodels.WorkflowStep{ + ID: "step1", WorkflowID: "wf1", Name: "Step 1", Position: 0, + } + taskRepo := newMockTaskRepo() + seedMockTaskState(taskRepo, "t1", v1.TaskStateInProgress) + agentMgr := &mockAgentManager{isAgentRunning: true} + agentMgr.currentPromptExecutionID = "exec-1" + agentMgr.currentPromptGeneration.Store(7) + svc := createTestServiceWithAgent(repo, stepGetter, taskRepo, agentMgr) + svc.turnService = &repoTurnService{repo: repo} + eb := &recordingEventBus{} + svc.eventBus = eb + + turn, err := svc.turnService.StartTurn(ctx, "s1") + require.NoError(t, err) + + // The ready event closes the turn synchronously, exactly as it does for + // the non-terminal case — this fires unconditionally on every successful + // completion, terminal or not. + svc.handleAgentReady(ctx, watcher.AgentEventData{ + TaskID: "t1", SessionID: "s1", AgentExecutionID: "exec-1", PromptGeneration: 7, + }) + + // The process exits sometime after: agent.completed marks the execution + // terminal. Its live-turn snapshot finds nothing, since ready already + // closed the turn above. + svc.markExecutionCompleted("s1", "exec-1") + + svc.handleAgentStreamEvent(ctx, &lifecycle.AgentStreamEventPayload{ + TaskID: "t1", + SessionID: "s1", + ExecutionID: "exec-1", + Data: &lifecycle.AgentStreamEventData{ + Type: agentEventComplete, + PromptGeneration: 7, + Usage: &streams.PromptUsage{InputTokens: 10, OutputTokens: 5}, + }, + }) + + usageEvent := findPromptUsageEvent(t, eb) + require.NotNil(t, usageEvent, "expected a session_prompt_usage.updated event to be published") + require.Equal(t, turn.ID, usageEvent.TurnID, + "terminal completion after a prior ready event must carry the turn id from the ready mark, not NULL") +} + +func findAllPromptUsageEvents(t *testing.T, eb *recordingEventBus, sessionID string) []lifecycle.SessionPromptUsageEventPayload { + t.Helper() + var out []lifecycle.SessionPromptUsageEventPayload + for _, rec := range eb.events { + if rec.subject != events.BuildSessionPromptUsageSubject(sessionID) { + continue + } + payload, ok := rec.event.Data.(lifecycle.SessionPromptUsageEventPayload) + require.True(t, ok, "session_prompt_usage.updated event carried an unexpected payload type") + out = append(out, payload) + } + return out +} + +// TestUsageEventIDFor is the unit-level regression test for F2: the id must +// be a deterministic function of (session, execution, prompt generation), +// not a value minted fresh per call. +func TestUsageEventIDFor(t *testing.T) { + a := usageEventIDFor("s1", "exec-1", 3) + b := usageEventIDFor("s1", "exec-1", 3) + require.Equal(t, a, b, "same (session, execution, prompt generation) must derive the same id") + require.NotEmpty(t, a) + + require.NotEqual(t, a, usageEventIDFor("s1", "exec-1", 4), + "a different prompt generation must derive a different id") + require.NotEqual(t, a, usageEventIDFor("s1", "exec-2", 3), + "a different execution must derive a different id") + require.NotEqual(t, a, usageEventIDFor("s2", "exec-1", 3), + "a different session must derive a different id") + + // promptGeneration==0 means the completion carries no generation + // tracking at all (see claimPromptCompletion's early return in the + // lifecycle package). Deriving a fixed key there would collide across + // genuinely distinct turns on a generation-less transport and silently + // under-count cost, so it must keep falling back to a random id. + require.NotEqual(t, usageEventIDFor("s1", "exec-1", 0), usageEventIDFor("s1", "exec-1", 0), + "promptGeneration==0 must not derive a stable id") +} + +// TestPublishPromptUsage_RepublishedCompletionReusesUsageEventID is the +// integration-level regression test for F2: a random UsageEventID minted on +// every publish can only dedup literal redelivery of the identical +// *bus.Event, which neither event bus provides (see +// internal/events/bus/{memory,nats}.go). The real duplicate source is the +// SAME completion frame reaching publishPromptUsage twice — e.g. a +// reconnecting WS client replaying a buffered stream event, mirroring the +// "late terminal complete" scenario already covered elsewhere in this file +// (TestCompleteStreamFromCompletedExecutionSkipsDuplicateOfficeTeardown and +// neighbors). Both publishes must carry the SAME usage_event_id so the +// office cost subscriber's unique index — proven separately by +// TestPromptUsage_DuplicateUsageEventIDIsIdempotent in internal/office/service +// — rejects the second row as a duplicate rather than double-counting cost. +func TestPublishPromptUsage_RepublishedCompletionReusesUsageEventID(t *testing.T) { + ctx := context.Background() + repo := setupTestRepo(t) + seedSession(t, repo, "t1", "s1", "") + + svc := createTestService(repo, newMockStepGetter(), newMockTaskRepo()) + svc.turnService = &repoTurnService{repo: repo} + eb := &recordingEventBus{} + svc.eventBus = eb + + _, err := svc.turnService.StartTurn(ctx, "s1") + require.NoError(t, err) + svc.markExecutionCompleted("s1", "exec-1") + svc.completeTurnForSession(ctx, "s1") + + frame := func() *lifecycle.AgentStreamEventPayload { + return &lifecycle.AgentStreamEventPayload{ + TaskID: "t1", + SessionID: "s1", + ExecutionID: "exec-1", + Data: &lifecycle.AgentStreamEventData{ + Type: agentEventComplete, + PromptGeneration: 3, + Usage: &streams.PromptUsage{InputTokens: 10, OutputTokens: 5}, + }, + } + } + + // Simulates the same buffered stream frame being delivered twice — same + // session, execution, and prompt generation both times. + svc.handleAgentStreamEvent(ctx, frame()) + svc.handleAgentStreamEvent(ctx, frame()) + + published := findAllPromptUsageEvents(t, eb, "s1") + require.Len(t, published, 2, "expected both republished frames to publish a prompt-usage event") + require.NotEmpty(t, published[0].UsageEventID) + require.Equal(t, published[0].UsageEventID, published[1].UsageEventID, + "republishing the same completion must reuse the same usage_event_id so the DB unique index catches the duplicate") +} + func TestCompleteStreamFromCompletedExecutionSkipsDuplicateOfficeTeardown(t *testing.T) { ctx := context.Background() repo := setupTestRepo(t) diff --git a/apps/backend/internal/orchestrator/executor/executor.go b/apps/backend/internal/orchestrator/executor/executor.go index 9cfe0dd06f..5266cd09f2 100644 --- a/apps/backend/internal/orchestrator/executor/executor.go +++ b/apps/backend/internal/orchestrator/executor/executor.go @@ -293,6 +293,13 @@ type AgentManagerClient interface { WaitForAgentctlReady(ctx context.Context, sessionID string) error } +// PromptTurnIDSetter is an optional lifecycle capability. Keeping it out of +// AgentManagerClient lets test and legacy adapters continue to work while the +// production lifecycle carries durable turn identity with completion events. +type PromptTurnIDSetter interface { + SetPromptTurnID(ctx context.Context, agentExecutionID, turnID string) error +} + // RemoteRuntimeStatus mirrors runtime status details needed by orchestrator/UI. type RemoteRuntimeStatus struct { RuntimeName agentruntime.Runtime @@ -336,6 +343,7 @@ type LaunchAgentRequest struct { TaskEnvironmentID string // Env owning this session (shared across sessions in the same task) TaskTitle string // Human-readable task title for semantic worktree naming AgentProfileID string + TurnID string // Durable Kandev turn for the initial prompt, when present // OfficeAgentProfileID is the stable Office identity. AgentProfileID stays // the concrete execution profile inside the executor for compatibility. OfficeAgentProfileID string @@ -477,6 +485,7 @@ type LaunchOptions struct { AgentProfileID string OfficeAgentProfileID string ExecutorID string + TurnID string Prompt string WorkflowStepID string StartAgent bool diff --git a/apps/backend/internal/orchestrator/executor/executor_execute.go b/apps/backend/internal/orchestrator/executor/executor_execute.go index b16a88808b..b5ab058084 100644 --- a/apps/backend/internal/orchestrator/executor/executor_execute.go +++ b/apps/backend/internal/orchestrator/executor/executor_execute.go @@ -975,7 +975,7 @@ func (e *Executor) LaunchPreparedSession(ctx context.Context, task *v1.Task, ses // overwrite the stale row. hasRunning, _ := e.repo.HasExecutorRunningRow(ctx, sessionID) if hasRunning { - result, err := e.startAgentOnExistingWorkspace(ctx, task, session, prompt, startAgent, opts.McpMode, opts.Env) + result, err := e.startAgentOnExistingWorkspace(ctx, task, session, prompt, startAgent, opts.McpMode, opts.Env, opts.TurnID) if !errors.Is(err, ErrStaleExecution) && !errors.Is(err, ErrAgentCommandMissing) { return result, err } @@ -1030,6 +1030,7 @@ func (e *Executor) LaunchPreparedSession(ctx context.Context, task *v1.Task, ses session.ExecutorID = execCfg.ExecutorID } req.OfficeAgentProfileID = opts.OfficeAgentProfileID + req.TurnID = opts.TurnID if req.OfficeAgentProfileID == "" && session.AgentProfileID != "" { req.OfficeAgentProfileID = session.AgentProfileID } @@ -1587,7 +1588,7 @@ func (e *Executor) applyRepositoryConfig(req *LaunchAgentRequest, task *v1.Task, // re-launch path. Pre-refactor this also consulted session.AgentExecutionID and // reconciled DB drift; that's now structurally impossible because executors_running // is owned by the lifecycle manager and writes are atomic with executionStore.Add. -func (e *Executor) startAgentOnExistingWorkspace(ctx context.Context, task *v1.Task, session *models.TaskSession, prompt string, startAgent bool, mcpMode string, env map[string]string) (*TaskExecution, error) { +func (e *Executor) startAgentOnExistingWorkspace(ctx context.Context, task *v1.Task, session *models.TaskSession, prompt string, startAgent bool, mcpMode string, env map[string]string, turnIDs ...string) (*TaskExecution, error) { executionID, err := e.agentManager.GetExecutionIDForSession(ctx, session.ID) if err != nil || executionID == "" { // No execution exists in memory (e.g. backend restarted since workspace was prepared). @@ -1624,32 +1625,9 @@ func (e *Executor) startAgentOnExistingWorkspace(ctx context.Context, task *v1.T // Non-fatal: agent may start without description } } - credentialReq := &LaunchAgentRequest{WorkspaceID: task.WorkspaceID, Env: cloneStringMap(env)} - e.injectGitLabWorkspaceCredentials(ctx, credentialReq) - if len(credentialReq.Env) > 0 { - if err := e.agentManager.SetExecutionEnv(ctx, executionID, credentialReq.Env); err != nil { - e.logger.Warn("failed to set execution env for existing workspace", - zap.String("session_id", session.ID), - zap.String("agent_execution_id", executionID), - zap.Error(err)) - } - } - - // If config MCP mode is needed, reconfigure the MCP server before starting the agent. - // The workspace may have been prepared before config_mode was set on the session. - effectiveMcpMode := mcpMode - if effectiveMcpMode == "" && isConfigModeSession(session) { - effectiveMcpMode = McpModeConfig - } - if effectiveMcpMode != "" { - if err := e.agentManager.SetMcpMode(ctx, executionID, effectiveMcpMode); err != nil { - e.logger.Error("failed to set MCP mode for existing workspace", - zap.String("session_id", session.ID), - zap.String("agent_execution_id", executionID), - zap.String("mcp_mode", effectiveMcpMode), - zap.Error(err)) - return nil, fmt.Errorf("set MCP mode %q: %w", effectiveMcpMode, err) - } + e.bindPromptTurnID(ctx, session.ID, executionID, turnIDs) + if err := e.configureExistingWorkspace(ctx, task, session, executionID, mcpMode, env); err != nil { + return nil, err } // Lazy workspace restoration creates an execution without an agent command. @@ -1695,6 +1673,62 @@ func (e *Executor) startAgentOnExistingWorkspace(ctx context.Context, task *v1.T return execution, nil } +func (e *Executor) configureExistingWorkspace( + ctx context.Context, + task *v1.Task, + session *models.TaskSession, + executionID, mcpMode string, + env map[string]string, +) error { + credentialReq := &LaunchAgentRequest{WorkspaceID: task.WorkspaceID, Env: cloneStringMap(env)} + e.injectGitLabWorkspaceCredentials(ctx, credentialReq) + if len(credentialReq.Env) > 0 { + if err := e.agentManager.SetExecutionEnv(ctx, executionID, credentialReq.Env); err != nil { + e.logger.Warn("failed to set execution env for existing workspace", + zap.String("session_id", session.ID), + zap.String("agent_execution_id", executionID), + zap.Error(err)) + } + } + + // If config MCP mode is needed, reconfigure the MCP server before starting the agent. + // The workspace may have been prepared before config_mode was set on the session. + effectiveMcpMode := mcpMode + if effectiveMcpMode == "" && isConfigModeSession(session) { + effectiveMcpMode = McpModeConfig + } + if effectiveMcpMode == "" { + return nil + } + if err := e.agentManager.SetMcpMode(ctx, executionID, effectiveMcpMode); err != nil { + e.logger.Error("failed to set MCP mode for existing workspace", + zap.String("session_id", session.ID), + zap.String("agent_execution_id", executionID), + zap.String("mcp_mode", effectiveMcpMode), + zap.Error(err)) + return fmt.Errorf("set MCP mode %q: %w", effectiveMcpMode, err) + } + return nil +} + +func (e *Executor) bindPromptTurnID(ctx context.Context, sessionID, executionID string, turnIDs []string) { + if len(turnIDs) == 0 || turnIDs[0] == "" { + return + } + setter, ok := e.agentManager.(PromptTurnIDSetter) + if !ok { + return + } + turnID := turnIDs[0] + if err := setter.SetPromptTurnID(ctx, executionID, turnID); err != nil { + e.logger.Warn("failed to bind prompt turn for existing workspace", + zap.String("session_id", sessionID), + zap.String("agent_execution_id", executionID), + zap.String("turn_id", turnID), + zap.Error(err)) + } +} + // captureBaseCommit retrieves the merge-base commit from agentctl and stores it // as the base commit for the session. This allows calculating cumulative diffs // that show all changes on the branch relative to the target branch (e.g., main). diff --git a/apps/backend/internal/orchestrator/service.go b/apps/backend/internal/orchestrator/service.go index e7a7ba0b8b..7849a13c3c 100644 --- a/apps/backend/internal/orchestrator/service.go +++ b/apps/backend/internal/orchestrator/service.go @@ -645,6 +645,30 @@ type Service struct { // guard does not grow without bound in long-running backend processes. completedExecutions sync.Map + // readyTurnMarks records, per (session, execution, prompt generation), + // the turn ID handleAgentReady confirmed and is about to close via + // completeTurnForSession. handleCompleteStreamEvent reads this snapshot + // (on both the terminal and non-terminal path) instead of trusting live + // active-turn state or the terminal-execution marker: agent.ready is + // published before the complete-stream frame for the same completion + // (see markReadyTurn's doc comment for the ordering guarantee and its + // NATS-deployment caveat), so both of those alternatives can already + // find the turn gone or stale by the time this runs. Entries are + // consumed on read and expire after the same grace window as + // completedExecutions so an unread entry cannot grow the map unbounded. + // This map is per-process — see markReadyTurn's doc comment for why a + // horizontally-scaled NATS deployment can miss cross-instance. + readyTurnMarks sync.Map + + // readyTurnMarksZeroGen is readyTurnMarks' sibling for promptGeneration==0 + // completions (transports with no generation tracking at all), which share + // one key per (session, execution) with no generation to disambiguate + // them — so entries queue FIFO here instead of occupying a single slot in + // readyTurnMarks. Guarded by readyTurnMarksZeroGenMu since sync.Map has no + // atomic append. See markReadyTurn's doc comment. + readyTurnMarksZeroGenMu sync.Mutex + readyTurnMarksZeroGen map[string][]readyTurnMark + // executionTeardownClaims arbitrates detached runtime teardown by // "::". Coordinator stop requests graceful // teardown while terminal-event cleanup requests force teardown; the first diff --git a/apps/backend/internal/orchestrator/task_operations.go b/apps/backend/internal/orchestrator/task_operations.go index 9283ac3823..5017826afc 100644 --- a/apps/backend/internal/orchestrator/task_operations.go +++ b/apps/backend/internal/orchestrator/task_operations.go @@ -547,10 +547,14 @@ func (s *Service) StartCreatedSession( if isOfficeTask { mcpMode = executor.McpModeOffice } - execution, err := s.executor.LaunchPreparedSession(ctx, task, sessionID, executor.LaunchOptions{AgentProfileID: effectiveProfileID, ExecutorID: executorID, Prompt: effectivePrompt, StartAgent: true, McpMode: mcpMode, Attachments: attachments}) + initialTurnID, initialTurnCreated := s.startTurnForSessionWithOwnership(ctx, sessionID) + execution, err := s.executor.LaunchPreparedSession(ctx, task, sessionID, executor.LaunchOptions{AgentProfileID: effectiveProfileID, ExecutorID: executorID, Prompt: effectivePrompt, StartAgent: true, McpMode: mcpMode, Attachments: attachments, TurnID: initialTurnID}) if err != nil { // The executor persists LaunchAgent failures. Cover earlier prepared-session // failures here; the session-level claim makes either completion order safe. + if initialTurnCreated { + s.completeTurnIfCurrent(ctx, sessionID, initialTurnID) + } return nil, s.handleSessionLaunchFailure(ctx, taskID, sessionID, err) } @@ -1014,6 +1018,7 @@ func (s *Service) startTask(ctx context.Context, taskID string, agentProfileID s if isOfficeTask { mcpMode = executor.McpModeOffice } + initialTurnID, initialTurnCreated := s.startTurnForSessionWithOwnership(ctx, sessionID) // Cache the raw prompt so a transient-provider-error (529) retry can // re-drive this first turn — initial launches bypass PromptTask. @@ -1023,6 +1028,7 @@ func (s *Service) startTask(ctx context.Context, taskID string, agentProfileID s AgentProfileID: agentProfileID, OfficeAgentProfileID: officeAgentProfileID, ExecutorID: executorID, + TurnID: initialTurnID, Prompt: effectivePrompt, WorkflowStepID: workflowStepID, StartAgent: true, @@ -1032,6 +1038,9 @@ func (s *Service) startTask(ctx context.Context, taskID string, agentProfileID s RouteOverride: route, }) if err != nil { + if initialTurnCreated { + s.completeTurnIfCurrent(ctx, sessionID, initialTurnID) + } return nil, s.handleSessionLaunchFailure(ctx, taskID, sessionID, err) } @@ -3569,6 +3578,16 @@ func (s *Service) promptTask(ctx context.Context, taskID, sessionID string, prom // Prompts can take a long time (minutes) while the WS request may timeout in 15 seconds. // We still want to log and respond, but the prompt should continue regardless. promptCtx := context.WithoutCancel(ctx) + if rollback.turnID != "" { + if setter, ok := s.agentManager.(executor.PromptTurnIDSetter); ok { + if err := setter.SetPromptTurnID(promptCtx, session.AgentExecutionID, rollback.turnID); err != nil { + s.logger.Warn("failed to bind prompt turn before dispatch", + zap.String("session_id", sessionID), + zap.String("turn_id", rollback.turnID), + zap.Error(err)) + } + } + } result, err := s.executor.PromptWithDispatchCallback( promptCtx, taskID, sessionID, effectivePrompt, attachments, dispatchOnly, func() { diff --git a/apps/backend/internal/task/repository/sqlite/session.go b/apps/backend/internal/task/repository/sqlite/session.go index 50cd5a81a0..29184ae874 100644 --- a/apps/backend/internal/task/repository/sqlite/session.go +++ b/apps/backend/internal/task/repository/sqlite/session.go @@ -1599,13 +1599,37 @@ func (r *Repository) GetLastAgentMessage(ctx context.Context, sessionID string) // stay in sync without re-summing office_cost_events. The model + DTO // don't surface these columns yet (DB-only per the office-costs // wedge); the cost explorer follow-up will expose them. +// +// Delegates to IncrementTaskSessionUsageTx using r.db as the executor; a +// caller that needs this atomic with another write (e.g. the office cost +// subscriber's ledger insert) should call the Tx variant directly with a +// shared transaction instead. func (r *Repository) IncrementTaskSessionUsage( ctx context.Context, sessionID string, tokensIn, tokensCachedIn, tokensOut, costSubcents int64, +) error { + return r.IncrementTaskSessionUsageTx(ctx, nil, sessionID, tokensIn, tokensCachedIn, tokensOut, costSubcents) +} + +// IncrementTaskSessionUsageTx implements shared.SessionUsageWriterTx: same +// write as IncrementTaskSessionUsage, but executed against tx when non-nil +// (falling back to r.db, the shared writer connection, when tx is nil) so a +// caller can make this atomic with another write in the same transaction. +func (r *Repository) IncrementTaskSessionUsageTx( + ctx context.Context, tx *sqlx.Tx, sessionID string, tokensIn, tokensCachedIn, tokensOut, costSubcents int64, ) error { if sessionID == "" { return nil } - _, err := r.db.ExecContext(ctx, r.db.Rebind(` + var exec interface { + ExecContext(ctx context.Context, query string, args ...interface{}) (sql.Result, error) + Rebind(query string) string + } + if tx != nil { + exec = tx + } else { + exec = r.db + } + _, err := exec.ExecContext(ctx, exec.Rebind(` UPDATE task_sessions SET tokens_in = COALESCE(tokens_in, 0) + ?, tokens_cached_in = COALESCE(tokens_cached_in, 0) + ?, diff --git a/docs/specs/office/costs.md b/docs/specs/office/costs.md index 1a26eac02f..82f2a40f4d 100644 --- a/docs/specs/office/costs.md +++ b/docs/specs/office/costs.md @@ -1,7 +1,7 @@ --- status: in-progress created: 2026-04-25 -updated: 2026-08-08 +updated: 2026-08-14 owner: cfl --- @@ -42,16 +42,26 @@ Office adds per-session cost estimation, aggregated views by agent/project/model | `model` | string | e.g. `claude-sonnet-4-20250514` | | `provider` | string | e.g. `anthropic`, `openai` | | `tokens_in` | int64 | input token count | -| `tokens_cached_in` | int64 | cached input tokens (prompt caching) | +| `tokens_cached_in` | int64 | cached input tokens (prompt caching). Always `= tokens_cached_read + tokens_cached_write` when the split is known; kept for existing rollup consumers (see `TaskSession additions` below) | +| `tokens_cached_read` | int64, nullable | cached-read portion of `tokens_cached_in`. NULL for legacy rows (pre-contract-v2) and for adapters with no per-turn usage frame (`estimated=true`, e.g. codex-acp) - never `0`, since a merged column cannot be decomposed after the fact | +| `tokens_cached_write` | int64, nullable | cached-write portion of `tokens_cached_in`. Same NULL rule as `tokens_cached_read` - cache write and cache read price at different per-million rates, so the split matters | | `tokens_out` | int64 | output token count | | `cost_subcents` | int64 | hundredths of a cent (renamed from `cost_cents`). UI divides by 10000 | -| `estimated` | bool | true when token counts were synthesised | +| `estimated` | bool | true when token counts were synthesised (e.g. codex-acp's cumulative-delta inference). A token-synthesis flag ONLY - independent of `cost_source` below. An unpriced row (list-price lookup miss) does not force this true; it reflects whatever the usage payload reported | +| `turn_id` | string, nullable | the turn this event was recorded at completion of. NULL when no turn could be resolved (e.g. an error/interrupt completion outside the normal ready path) | +| `usage_event_id` | string, nullable | idempotency key derived from `(session_id, execution_id, prompt_generation)` at the publish site; unique when non-NULL (partial index) so a redelivered prompt-usage event does not double-insert or double-count the session rollup. NULL for legacy rows. Generation-less transports receive a random per-publish key, so those rows are intentionally not deduplicated | +| `cost_source` | enum, nullable | `provider_reported` \| `models_dev_list` \| `unpriced`. Records which layer of the two-layer lookup below actually produced `cost_subcents` - distinct from `estimated`. NULL for legacy rows | +| `rate_input_per_million` / `rate_cached_read_per_million` / `rate_cached_write_per_million` / `rate_output_per_million` | int64, nullable | the four models.dev rates actually applied, in subcents-per-million-tokens. Populated only when `cost_source = models_dev_list`; NULL otherwise (including `provider_reported`, where no rate table applies) | +| `pricing_catalog_version` | string, nullable | the models.dev cache's "as-of" identifier (RFC3339 load/fetch time - the dataset has no version field of its own) that produced the applied rates. Populated only when `cost_source = models_dev_list` | +| `cost_contract_version` | int64, nullable | in-band activation marker for this row's column semantics, since the point-in-time Rill extract has no schema versioning of its own. Current value `2` (v1 -> v2: the unpriced path stopped forcing `estimated=true`, see the `estimated` row above). NULL for legacy pre-contract rows - never backfilled | | `provider_event_id` | string | null for wire-side rows; non-null for disk-runner rows. Format `ccusage:::` | | `provider_credits` | int64 | amp `credits` field (null elsewhere) | | `occurred_at` | timestamp | when the cost was incurred | Disk-runner rows are upserted by `(session_id, provider_event_id)`. For codex, after a successful aggregate row lands, the wire-side `estimated=true` rows for the same session are deleted to avoid double-counting. +**Contract v2 columns** (`tokens_cached_read`, `tokens_cached_write`, `turn_id`, `usage_event_id`, `cost_source`, the four `rate_*_per_million` columns, `pricing_catalog_version`, `cost_contract_version`) were added via an idempotent nullable-column migration (`migrateCostEventContract`); every legacy row reads NULL for all of them, never `0` - see the cross-cutting rule in `Persistence guarantees` below. Postgres parity follows the same migration shape per ADR 0027. + ### `office_budget_policies` | Field | Type | Notes | @@ -68,7 +78,7 @@ Multiple policies can apply to the same scope (e.g. monthly + total). ### `TaskSession` additions -`cost_subcents`, `tokens_in`, `tokens_out`, `tokens_cached_in` incrementally updated as cost events arrive, providing quick per-session totals without scanning `office_cost_events`. `tokens_cached_in` is DB-only today (no model/DTO/API/web reader) — it exists so the rollup reconciles column-for-column against the ledger; a future cost-explorer surface will expose it. +`cost_subcents`, `tokens_in`, `tokens_out`, `tokens_cached_in` incrementally updated as cost events arrive, providing quick per-session totals without scanning `office_cost_events`. `tokens_cached_in` is DB-only today (no model/DTO/API/web reader) - it exists so the rollup reconciles column-for-column against the ledger; a future cost-explorer surface will expose it. The rollup consumes only the merged `tokens_cached_in` sum, not the `tokens_cached_read` / `tokens_cached_write` split - the split exists solely on the ledger row for cost-per-token-type analysis (contract v2); it does not change what the rollup total means or how it is computed. ### Per-agent budget @@ -106,10 +116,10 @@ replaced successfully. ### Cost resolution (two-layer lookup) -1. **Provider-reported cost.** If the CLI emits cumulative USD session cost on `usage_update`, subtract the previously consumed session baseline and store the nonnegative turn delta as `int64(amount * 10000)` hundredths-of-a-cent. Ignore non-USD cost values. This is the only accurate path for claude-acp. -2. **`models.dev` lookup.** For CLIs reporting tokens but no cost (gemini, opencode BYOK, codex fallback), resolve pricing against the cached dataset. +1. **Provider-reported cost** (`cost_source = provider_reported`). If the CLI emits cumulative USD session cost on `usage_update`, subtract the previously consumed session baseline and store the nonnegative turn delta as `int64(amount * 10000)` hundredths-of-a-cent. Ignore non-USD cost values. This is the only accurate path for claude-acp. No rate columns are recorded on this path. +2. **`models.dev` lookup** (`cost_source = models_dev_list`). For CLIs reporting tokens but no cost (gemini, opencode BYOK, codex fallback), resolve pricing against the cached dataset. The four applied per-million rates and the catalogue's `pricing_catalog_version` are recorded alongside `cost_subcents`. -When both miss (first-boot, no network, model unknown, proprietary id), the row records `cost_subcents=0` and `estimated=true`. UI shows "pricing unavailable". Users can override pricing per model in workspace settings. +When both miss (first-boot, no network, model unknown, proprietary id), the row records `cost_subcents=0` and `cost_source=unpriced`. `estimated` is NOT forced true on this path - it independently tracks whether the token counts themselves were synthesised (see the `office_cost_events` table above). Before contract v2, an unpriced row was indistinguishable from a token-synthesis row; downstream consumers that read `estimated` as a stand-in for pricing confidence should switch to `cost_source`. UI shows "pricing unavailable". Users can override pricing per model in workspace settings. ### Disk-runner binary: `cmd/usage-runner` @@ -169,7 +179,10 @@ There is no per-field permission model. Conformance tests should assert that cos ## Failure modes -- **models.dev miss**: row recorded with `cost_subcents=0` and `estimated=true`; UI shows "pricing unavailable". +- **models.dev miss**: row recorded with `cost_subcents=0` and `cost_source=unpriced`; `estimated` is unaffected (tracks token-synthesis independently, see the `office_cost_events` table). UI shows "pricing unavailable". +- **Prompt-usage event redelivered** (e.g. a reconnecting WS client replaying a buffered stream event): the second `CreateCostEvent` call collides on the unique `usage_event_id` partial index and is dropped as a no-op rather than double-inserting or double-counting the session rollup. +- **Writer silently stops**: every path a prompt-usage event can take (decode failure, missing task/session ids, task-fields lookup failure, insert failure or duplicate, or a successful write) increments an `expvar` counter (`cost_events_written_total` / `cost_events_dropped_total`, dev-mode `/debug/vars`) and logs, so a stopped writer is observable instead of silently producing zero rows. +- **Ledger insert and rollup increment commit or roll back together**: `CreateCostEvent` (the `office_cost_events` insert) and `IncrementTaskSessionUsage` (the `task_sessions` rollup) are wrapped in one `*sqlx.Tx` by `recordCostEventAndRollup` whenever the wired session-usage writer supports it (every production wiring does). If the rollup increment fails after the ledger insert, the whole transaction rolls back - nothing is committed - so a redelivery of the same event with the same `usage_event_id` retries the operation cleanly instead of being dropped as a duplicate by the unique index above while the rollup stays short. Falls back to the pre-existing non-atomic two-step write only when the wired writer doesn't implement the transactional capability (for example a simplified test double). - **models.dev fetch fails or is canceled**: background refresh falls back to existing on-disk and in-memory data; no crash. The in-flight guard is released and a later refresh can retry. @@ -185,9 +198,10 @@ There is no per-field permission model. Conformance tests should assert that cos **Survives restart:** -- `office_cost_events` rows (full history, never trimmed). Disk-runner rows keyed by `(session_id, provider_event_id)` survive re-ingestion without duplicating; the wire-side `estimated=true` rows for codex are deleted once the matching aggregate row lands. +- `office_cost_events` rows (full history, never trimmed). Disk-runner rows keyed by `(session_id, provider_event_id)` survive re-ingestion without duplicating; wire-side rows are additionally deduplicated by the unique `usage_event_id` partial index (contract v2). The wire-side `estimated=true` rows for codex are deleted once the matching aggregate row lands. +- **Contract v2 column semantics never change value mid-series.** Legacy rows (written before contract v2) read NULL, never `0`, for every new column - the downstream point-in-time extract has no schema versioning, so a column whose meaning flips from "not recorded" to "recorded as zero" partway through the series would be silently discontinuous. `cost_contract_version` on each row is the in-band marker a consumer checks instead of inferring from a date. - `office_budget_policies` rows. -- `TaskSession.cost_subcents` / `tokens_in` / `tokens_out` / `tokens_cached_in` running totals, kept in sync with `office_cost_events` so per-session totals are correct without a re-scan. +- `TaskSession.cost_subcents` / `tokens_in` / `tokens_out` / `tokens_cached_in` running totals, kept in sync with `office_cost_events` so per-session totals are correct without a re-scan - written atomically with the ledger insert (see `Failure modes` above), so the two can never diverge on a partial write. - Per-agent `budget_monthly_cents` stored on the agent instance row. - Per-model pricing overrides stored in workspace settings. - The on-disk `models.dev` cache at `/cache/models-dev.json`. Recovery on next boot: the in-memory pricing map is empty until first query; queries fall back to the on-disk file when the background refresh has not yet completed. @@ -209,7 +223,19 @@ No TTL or retention is applied to `office_cost_events`; rows accumulate for the - **GIVEN** an ACP agent session processing a turn, **WHEN** its `complete` event carries normalized usage, **THEN** a cost event is created from the provider-reported USD delta or estimated via models.dev pricing. The session's cumulative `cost_subcents` is updated. -- **GIVEN** a model not found in the models.dev dataset and no user override configured, **WHEN** a cost event is recorded, **THEN** token counts are stored but `cost_subcents` is zero. The cost explorer shows "pricing unavailable" for that model. +- **GIVEN** a model not found in the models.dev dataset and no user override configured, **WHEN** a cost event is recorded, **THEN** token counts are stored but `cost_subcents` is zero and `cost_source=unpriced`. The cost explorer shows "pricing unavailable" for that model. + +- **GIVEN** a claude-acp turn that reports both cached-read and cached-write tokens, **WHEN** the cost event is recorded, **THEN** `tokens_cached_read` and `tokens_cached_write` are stored as the separate values reported (not merged), `tokens_cached_in` equals their sum, and `cost_subcents` reflects each portion priced at its own per-million rate. + +- **GIVEN** a codex-acp turn (no per-turn usage frame; tokens synthesised from context-occupancy growth, `estimated=true`), **WHEN** the cost event is recorded, **THEN** `tokens_cached_read` and `tokens_cached_write` are NULL (never `0`) because the split was never observed, while `tokens_cached_in` and `cost_subcents` are still recorded from the synthesised totals. + +- **GIVEN** a turn priced via the models.dev lookup, **WHEN** the cost event is recorded, **THEN** `cost_source=models_dev_list` and the row carries the four `rate_*_per_million` values and `pricing_catalog_version` actually applied, distinguishable from a provider-reported row (`cost_source=provider_reported`, no rate columns) and an unpriced row (`cost_source=unpriced`, `cost_subcents=0`). + +- **GIVEN** a prompt-usage event for an in-progress turn, **WHEN** the cost event is recorded at turn completion, **THEN** `turn_id` is populated with the turn that was resolved at completion time and `usage_event_id` is set to a deterministic id derived from the underlying completion's identity. + +- **GIVEN** the same underlying prompt-usage completion is delivered twice (e.g. a reconnecting WS client replaying a buffered stream event), **WHEN** the second copy is processed, **THEN** it derives the same `usage_event_id`, the insert is rejected by the unique partial index, and the session rollup is not incremented a second time. + +- **GIVEN** an `office_cost_events` table that predates contract v2, **WHEN** the repository boots and runs its migration, **THEN** the new columns are added nullable, every pre-existing row reads NULL (never `0`) for all of them, and running the migration again on the same database is a no-op. - **GIVEN** many stale pricing lookups and direct refresh calls arrive together, **WHEN** refresh starts, **THEN** one network fetch runs for that client while diff --git a/docs/specs/workspaces/worktree-branch-templates.md b/docs/specs/workspaces/worktree-branch-templates.md index a84c6fa587..8d145d5d36 100644 --- a/docs/specs/workspaces/worktree-branch-templates.md +++ b/docs/specs/workspaces/worktree-branch-templates.md @@ -61,4 +61,3 @@ a non-empty template. - Recovery of a custom template that an earlier Kandev version already replaced. - Changes to the template syntax, placeholders, validation, default value, or UI. - Replacement of the current schema migration system. -