From 99850f1dc5c95434cdeb116adb3f013ef23453a2 Mon Sep 17 00:00:00 2001 From: Quanzheng Long Date: Thu, 24 Sep 2026 00:01:44 -0700 Subject: [PATCH] agent: add optional model usage hooks --- ARCHITECTURE.md | 15 +- README.md | 25 + agent/agent.go | 9 + ...17-optional-model-hooks-and-exact-usage.md | 49 ++ docs/flow-model.md | 17 + docs/v0.7.0-upgrade.md | 15 + internal/agent/flow.go | 619 ++++++++++++++++-- internal/agent/flow_integration_test.go | 306 ++++++++- internal/agent/model_hooks_test.go | 73 +++ internal/agent/types.go | 27 +- internal/app/app.go | 1 + internal/config/config.go | 65 +- internal/config/config_test.go | 31 +- internal/model/anthropic.go | 14 +- internal/model/client.go | 12 +- internal/model/gemini.go | 14 +- internal/model/groq.go | 14 +- internal/model/mock.go | 38 +- internal/model/mock_test.go | 15 +- internal/model/openai.go | 101 +-- internal/model/openai_live_test.go | 5 +- internal/model/openai_test.go | 72 +- internal/model/providers_test.go | 2 +- .../public-api-consumer/consumer_test.go | 8 +- 24 files changed, 1384 insertions(+), 163 deletions(-) create mode 100644 docs/adr/0017-optional-model-hooks-and-exact-usage.md create mode 100644 docs/v0.7.0-upgrade.md create mode 100644 internal/agent/model_hooks_test.go diff --git a/ARCHITECTURE.md b/ARCHITECTURE.md index a3e37cd..9b3c56e 100644 --- a/ARCHITECTURE.md +++ b/ARCHITECTURE.md @@ -60,7 +60,12 @@ depending on `internal/model` or duplicating provider wiring. `model.OpenAIClientConfig` controls whether the OpenAI Responses adapter asks for and forwards provider-authored reasoning summaries. It is disabled by default; encrypted reasoning content remains included for stateless turn -continuation regardless of that display setting. +continuation regardless of that display setting or response-storage policy. + +The same configuration independently controls response storage through +`DisableModelResponseStore`. Every OpenAI Complete and Summarize request sends +`store` explicitly. Storage defaults to enabled and is disabled only when that +field is true. ## Durable Agent model @@ -154,6 +159,14 @@ boundary. Waiting state is written in the `WaitFor` that establishes the wait. Provider and MCP calls occur only in `Execute`. The complete graph and resource table are in `docs/flow-model.md`. +`WithModelHooks` adds four durable model boundaries without changing the +legacy Step identities. Before admission generates a stable call ID from the +Flow and Step execution. The Call Step performs the provider request and +overwrites the single `ModelUsage` Attribute. The After Step reads that +Attribute and invokes the idempotent application hook; Apply then mutates Agent +history. No accounting totals, provider callback, model checkpoint state, or +model-recovery API exists in SuperAgent. + History-reading Steps declare bounded AttributeMap loads explicitly. Integration tests use verb-first `ForTestOnly` RPCs rather than generic Dex Client resource operations. Tool invocations receive the stable Flow ID, model call ID, and diff --git a/README.md b/README.md index b671003..3a45c39 100644 --- a/README.md +++ b/README.md @@ -79,6 +79,25 @@ idempotent application cleanup request. Accepted user mutations coalesce Timer resets within one minute. A `durable_wait` extends the deadline through its target; model and ordinary tool work do not pause it. +Embedders that need admission and accounting boundaries can pass +`agent.WithModelHooks(...)`. A non-nil configuration requires both hooks. Dex +runs `BeforeModelCall`, `CallModelWithHooks`, `AfterModelCall`, and +`ApplyModelResult` as separate Steps, so an After retry never repeats the model +call. The Call Step overwrites the `ModelUsage` Attribute with the provider's +exact usage for that call; the After hook receives that same value. SuperAgent +does not accumulate account usage. + +```go +flow := agent.NewFlow(modelClient, tools, agent.WithModelHooks(&agent.ModelHooksConfig{ + BeforeModelCall: admitModelCall, + AfterModelCall: synchronizeModelUsage, +})) +``` + +The Before hook returns `false, nil` for a business rejection. Hook errors use +the owning Dex Step retry policy. The After hook must be safely repeatable by +its supplied `CallID`. + External tool retries belong to Dex. `ToolDefinition` controls attempt timeout, maximum attempts, and total duration. A registry performs exactly one call for each Dex attempt. @@ -152,6 +171,7 @@ available. | `DEX_BLOB_CACHE_MAX_BYTES` | BlobCache size limit in bytes | `536870912` | | `DEX_AGENT_MCP_CONFIG` | Trusted MCP YAML path | disabled | | `OPENAI_API_KEY` | OpenAI credential | unset | +| `OPENAI_DISABLE_MODEL_RESPONSE_STORE` | Send `store=false` for OpenAI Responses calls | `false` | | `ANTHROPIC_API_KEY` | Anthropic credential | unset | | `GEMINI_API_KEY` | Gemini credential | unset | | `GROQ_API_KEY` | Groq credential | unset | @@ -162,6 +182,11 @@ persisted in Dex state or logged. Copy [`web/mcp-servers.example.yaml`](web/mcp-servers.example.yaml) to configure trusted MCP servers. +OpenAI Responses calls always send `store` explicitly. The default is +`store=true`; set `OPENAI_DISABLE_MODEL_RESPONSE_STORE=true` only when the +deployment requires provider-side response storage to be disabled. This setting +is independent of model hooks. + Each configured tool defaults to `running_type: short_running`. Use `long_running` when more than half of expected calls exceed five seconds. This is a Dex placement optimization, not a timeout or SLA; short-running calls may diff --git a/agent/agent.go b/agent/agent.go index 84d0fc1..868d15d 100644 --- a/agent/agent.go +++ b/agent/agent.go @@ -232,6 +232,7 @@ type ( InputConsumption = agentinternal.InputConsumption StreamEvent = agentinternal.StreamEvent ModelReply = agentinternal.ModelReply + ModelUsage = agentinternal.ModelUsage ToolDefinition = agentinternal.ToolDefinition ToolExecutionResult = agentinternal.ToolExecutionResult TextWriter = agentinternal.TextWriter @@ -239,6 +240,9 @@ type ( ModelRequest = agentinternal.ModelRequest SummarizeRequest = agentinternal.SummarizeRequest ModelClient = agentinternal.ModelClient + BeforeModelCallHook = agentinternal.BeforeModelCallHook + AfterModelCallHook = agentinternal.AfterModelCallHook + ModelHooksConfig = agentinternal.ModelHooksConfig ToolInvocation = agentinternal.ToolInvocation ToolRegistry = agentinternal.ToolRegistry RegisteredTool = agentinternal.RegisteredTool @@ -257,6 +261,11 @@ func WithInactivityExpirationHandler(handler InactivityExpirationHandler) FlowOp return agentinternal.WithInactivityExpirationHandler(handler) } +// WithModelHooks configures optional durable boundaries around model calls. +func WithModelHooks(config *ModelHooksConfig) FlowOption { + return agentinternal.WithModelHooks(config) +} + // NewClient constructs an Agent application client over one Dex client and Flow definition. func NewClient(sdkClient *dex.Client, flow *Flow) *Client { return agentinternal.NewClient(sdkClient, flow) diff --git a/docs/adr/0017-optional-model-hooks-and-exact-usage.md b/docs/adr/0017-optional-model-hooks-and-exact-usage.md new file mode 100644 index 0000000..350e625 --- /dev/null +++ b/docs/adr/0017-optional-model-hooks-and-exact-usage.md @@ -0,0 +1,49 @@ +# ADR 0017: Optional model hooks and exact per-call usage + +## Status + +Accepted for SuperAgent `v0.7.0`. + +## Context + +Embedding applications may need to admit model calls against account policy and +synchronize exact provider usage. That policy is application-owned. Putting +accounts, projects, quotas, totals, or recovery state in SuperAgent would couple +the reusable Agent Flow to one control plane. + +The model provider response is an external effect. Retrying accounting after a +successful provider call must not repeat that call, and retrying Agent state +application must not repeat accounting. + +## Decision + +`ModelClient.Complete` and `Summarize` return `ModelUsage` separately from their +existing result. `ModelReply` remains unchanged. Every successful legacy or +hooked Call Step overwrites one `ModelUsage` Attribute; it never accumulates +usage. + +`WithModelHooks` is optional and requires both a Before and After hook. Hooked +calls execute as four Dex Steps: + +1. `BeforeModelCall` creates a stable call ID and performs admission. +2. `CallModelWithHooks` calls the provider and writes `ModelUsage`. +3. `AfterModelCall` reads that Attribute and synchronizes usage idempotently. +4. `ApplyModelResult` applies the persisted provider result. + +The call ID derives from Flow ID, the Before Step execution ID, and call kind. +It is carried in Step input, not an Attribute. OpenAI receives it as its +idempotency key. Hooked calls permit only OpenAI and Mock until other adapters +report exact usage. + +OpenAI response storage is an independent client setting. Complete and +Summarize always send `store`; the default is true and +`DisableModelResponseStore` changes it to false. + +## Consequences + +After retries cannot repeat the provider call, and Apply retries cannot repeat +the After hook. The embedding application owns account lookup, monthly totals, +limits, and idempotency for synchronized usage. + +SuperAgent has no usage reporter, accumulated usage Attribute, model checkpoint +state, pending recovery record, model recovery API, or recovery UI. diff --git a/docs/flow-model.md b/docs/flow-model.md index 947eb50..7b93a83 100644 --- a/docs/flow-model.md +++ b/docs/flow-model.md @@ -35,6 +35,11 @@ retry, and recovery settings. Every regular attempt uses a one-minute heartbeat timeout. Ordinary Step methods use a one-minute method timeout, while model methods retain their explicit ten-minute method timeout. +When `WithModelHooks` is absent, `CompactContext` and `CallModel` keep their +existing identities and inputs. When it is present, both model-call kinds use +the four-Step hook path below. Only OpenAI and Mock provide exact usage in this +release, so a hooked Flow rejects other providers before dispatch. + The Worker negotiates the highest common protocol with the Server before Attribute index synchronization or Worker binding. Deploy the Server before the Worker. Startup fails when `GetServerInfo` is missing, either interval is @@ -64,6 +69,13 @@ CallModel -> CheckSteered -> AwaitUser (ordinary or repeated no-progress response) -> CheckSteered -> RouteTool (tool calls) +Hooks-enabled response or compaction + -> BeforeModelCall (admit and create stable CallID) + -> CallModelWithHooks (provider call; overwrite ModelUsage) + -> AfterModelCall (read ModelUsage; synchronize usage) + -> ApplyModelResult (mutate Agent history or summary) + -> existing continuation + RouteTool -> CheckSteered -> AwaitToolApproval (untrusted write) -> CheckSteered -> ExecuteTool (approved/read-only external tool) @@ -161,6 +173,10 @@ history, and makes the model replan. | `AnsweredInput` | none | Route an accepted answer through compaction directly to the model before checking steering | | `CompactContext` | none | Call the summary provider, commit the covered range and summary, then trim only summarized retained messages | | `CallModel` | none | Rebuild context, stream buffered deltas, commit the assistant message and pending calls; retry one active-plan response that made no durable progress | +| `BeforeModelCall` | none | Resolve the selected model, create the stable call ID, and invoke optional admission | +| `CallModelWithHooks` | none | Call the provider once per Dex attempt, overwrite exact `ModelUsage`, and persist the result in Step output | +| `AfterModelCall` | none | Read `ModelUsage` and invoke the application-owned idempotent synchronization hook | +| `ApplyModelResult` | none | Apply the persisted assistant response or compaction summary without repeating model or hook side effects | | `CheckSteered` | bounded steered batch | Apply steering at a safe boundary or route the explicit continuation | | `RouteTool` | none | Validate built-in arguments and select approval, MCP execution, timer, input, or next-call path | | `AwaitToolApproval` | exact call-ID approval or steering | Persist waiting status; consume one decision or replan on steering | @@ -228,6 +244,7 @@ deadline. `DurableWait` advances it to the wait target plus the timeout. | `AgentConfig` | Attribute | Immutable execution configuration | | `AgentRuntimeMetadata` | Attribute | Trusted runtime routing metadata; never model or browser context | | `AgentState` | Attribute | Sequence range, mode, status, plan revision, pending-call cursor, and Plan no-progress count | +| `ModelUsage` | Attribute | Exact usage from the last successful model call; overwritten, never accumulated | | `WaitingInputRound` | Attribute | Monotonic durable browser watermark, initialized to zero | | `ContextSummary` | Attribute | Cumulative summary and explicit covered sequence | | `CurrentMessages` | AttributeMap | Recent provider-neutral messages keyed by sequence | diff --git a/docs/v0.7.0-upgrade.md b/docs/v0.7.0-upgrade.md new file mode 100644 index 0000000..cd3d13c --- /dev/null +++ b/docs/v0.7.0-upgrade.md @@ -0,0 +1,15 @@ +# SuperAgent v0.7.0 upgrade + +`ModelClient` implementations must return one additional `ModelUsage` value +from `Complete` and `Summarize`. Existing `ModelReply` values do not change. +Adapters without exact usage may return `ModelUsage{}`; such providers cannot be +used with `WithModelHooks` in this release. + +Applications may opt into durable admission and usage synchronization with +`WithModelHooks`. Both hooks are required when the option is non-nil. After hooks +must be idempotent by the supplied stable call ID. + +OpenAI now explicitly sends `store=true` by default for both Complete and +Summarize. Set `OpenAIClientConfig.DisableModelResponseStore` or the standalone +process environment variable `OPENAI_DISABLE_MODEL_RESPONSE_STORE=true` to send +`store=false`. diff --git a/internal/agent/flow.go b/internal/agent/flow.go index 62e04a6..2a60fcc 100644 --- a/internal/agent/flow.go +++ b/internal/agent/flow.go @@ -26,6 +26,7 @@ import ( "strings" "time" + "github.com/google/uuid" "github.com/superdurable/dex/sdk-go/dex" ) @@ -33,6 +34,7 @@ var ( agentConfigAttribute = dex.DefineAttribute[AgentConfig]("AgentConfig") agentRuntimeMetadataAttribute = dex.DefineAttribute[JSONObject]("AgentRuntimeMetadata") agentStateAttribute = dex.DefineAttribute[AgentState]("AgentState") + modelUsageAttribute = dex.DefineAttribute[ModelUsage]("ModelUsage") waitingInputRoundAttribute = dex.DefineAttribute[WaitingInputRound]("WaitingInputRound") contextSummaryAttribute = dex.DefineAttribute[ContextSummary]("ContextSummary") currentMessagesAttribute = dex.DefineAttributeMap[AgentMessage]("CurrentMessages") @@ -59,6 +61,7 @@ var ( // Flow is the durable AI Agent state machine. type Flow struct { modelClient ModelClient + modelHooks *ModelHooksConfig tools ToolRegistry inactivityExpiration InactivityExpirationHandler rpcDefinitionsForTestOnly []dex.RPCDef @@ -81,6 +84,16 @@ func WithInactivityExpirationHandler(handler InactivityExpirationHandler) FlowOp } } +// WithModelHooks configures optional durable boundaries around model calls. +func WithModelHooks(config *ModelHooksConfig) FlowOption { + return func(flow *Flow) { + if config != nil && (config.BeforeModelCall == nil || config.AfterModelCall == nil) { + panic("both model hooks are required") + } + flow.modelHooks = config + } +} + // NewFlow constructs an Agent from its model and trusted tool boundaries. func NewFlow(modelClient ModelClient, tools ToolRegistry, options ...FlowOption) *Flow { if modelClient == nil { @@ -129,6 +142,10 @@ func (flow *Flow) GetSteps() []dex.StepDef { dex.DefineStep(answeredInputStep{flow: flow}), dex.DefineStep(compactContextStep{flow: flow}), dex.DefineStep(callModelStep{flow: flow}), + dex.DefineStep(beforeModelCallStep{flow: flow}), + dex.DefineStep(callModelWithHooksStep{flow: flow}), + dex.DefineStep(afterModelCallStep{flow: flow}), + dex.DefineStep(applyModelResultStep{flow: flow}), dex.DefineStep(checkSteeredStep{flow: flow}), dex.DefineStep(routeToolStep{flow: flow}), dex.DefineStep(awaitToolApprovalStep{flow: flow}), @@ -222,6 +239,7 @@ func (*Flow) GetPersistenceSchema() dex.PersistenceSchema { agentConfigAttribute, agentRuntimeMetadataAttribute, agentStateAttribute, + modelUsageAttribute, waitingInputRoundAttribute, contextSummaryAttribute, currentMessagesAttribute, @@ -1934,6 +1952,10 @@ const ( stepTypeAnsweredInput stepType = "AnsweredInput" stepTypeCompactContext stepType = "CompactContext" stepTypeCallModel stepType = "CallModel" + stepTypeBeforeModelCall stepType = "BeforeModelCall" + stepTypeCallModelWithHooks stepType = "CallModelWithHooks" + stepTypeAfterModelCall stepType = "AfterModelCall" + stepTypeApplyModelResult stepType = "ApplyModelResult" stepTypeCheckSteered stepType = "CheckSteered" stepTypeRouteTool stepType = "RouteTool" stepTypeAwaitApproval stepType = "AwaitToolApproval" @@ -1954,6 +1976,54 @@ const ( type continuation string type stepType string +type modelCallKind string + +const ( + modelCallKindResponse modelCallKind = "response" + modelCallKindCompaction modelCallKind = "compaction" +) + +type modelCallPreparation struct { + Kind modelCallKind `json:"kind"` + ThroughSequence Sequence `json:"through_sequence,omitempty"` +} + +type modelCallInput struct { + CallID CallID `json:"call_id"` + Kind modelCallKind `json:"kind"` + ThroughSequence Sequence `json:"through_sequence,omitempty"` +} + +type modelResponseResult struct { + Reply ModelReply `json:"reply"` + StartedAt time.Time `json:"started_at"` + ExpectedMessageSequence Sequence `json:"expected_message_sequence"` +} + +type modelCompactionResult struct { + Summary string `json:"summary"` + ThroughSequence Sequence `json:"through_sequence"` +} + +type modelCallResult struct { + CallID CallID `json:"call_id"` + Kind modelCallKind `json:"kind"` + Response *modelResponseResult `json:"response,omitempty"` + Compaction *modelCompactionResult `json:"compaction,omitempty"` +} + +func modelCallID(ctx dex.Context, kind modelCallKind) CallID { + encoded, err := json.Marshal(struct { + FlowID string `json:"flow_id"` + StepExecutionID string `json:"step_execution_id"` + Kind modelCallKind `json:"kind"` + }{ctx.FlowID(), ctx.StepExecutionID(), kind}) + if err != nil { + panic(fmt.Sprintf("encode model call identity: %v", err)) + } + return CallID(uuid.NewSHA1(uuid.NameSpaceOID, encoded).String()) +} + type heartbeatPhase string const ( @@ -2340,9 +2410,18 @@ func (step answeredInputStep) Execute(ctx dex.Context, _ dex.None) (*dex.StepDec return nil, err } if cutoff != nil { - return dex.GoTo(compactContextStep{flow: step.flow}, *cutoff), nil + if step.flow.modelHooks == nil { + return dex.GoTo(compactContextStep{flow: step.flow}, *cutoff), nil + } + return dex.GoTo(beforeModelCallStep{flow: step.flow}, modelCallPreparation{ + Kind: modelCallKindCompaction, + ThroughSequence: *cutoff, + }), nil + } + if step.flow.modelHooks == nil { + return dex.GoTo(callModelStep{flow: step.flow}, nil), nil } - return dex.GoTo(callModelStep{flow: step.flow}, nil), nil + return dex.GoTo(beforeModelCallStep{flow: step.flow}, modelCallPreparation{Kind: modelCallKindResponse}), nil } type compactContextStep struct { @@ -2368,12 +2447,7 @@ func (step compactContextStep) Execute(ctx dex.Context, input Sequence) (*dex.St if err != nil { return nil, err } - messages, err := step.flow.loadMessages( - ctx, - state.SummarizedThroughSequence+1, - input, - config, - ) + messages, err := step.flow.loadMessages(ctx, state.SummarizedThroughSequence+1, input, config) if err != nil { return nil, err } @@ -2382,28 +2456,28 @@ func (step compactContextStep) Execute(ctx dex.Context, input Sequence) (*dex.St return nil, err } if heartbeatErr := ctx.RecordHeartbeat(compactionHeartbeat{ - Phase: heartbeatPhaseCompacting, - ThroughSequence: input, + Phase: heartbeatPhaseCompacting, ThroughSequence: input, }); heartbeatErr != nil { return nil, heartbeatErr } - summary, err := step.flow.modelClient.Summarize(ctx, SummarizeRequest{ - Config: config, - PreviousSummary: previousSummary.Content, - Messages: messages, - FlowID: FlowID(ctx.FlowID()), + summary, usage, err := step.flow.modelClient.Summarize(ctx, SummarizeRequest{ + Config: config, PreviousSummary: previousSummary.Content, Messages: messages, + FlowID: FlowID(ctx.FlowID()), CallID: modelCallID(ctx, modelCallKindCompaction), }) if err != nil { - if eventErr := step.flow.writeActivity(ctx, AgentEvent{Kind: EventKindCompactionFailed, Message: "Context compaction failed."}); eventErr != nil { + if eventErr := step.flow.writeActivity(ctx, AgentEvent{ + Kind: EventKindCompactionFailed, Message: "Context compaction failed.", + }); eventErr != nil { return nil, errors.Join(err, eventErr) } return nil, err } + if usageErr := modelUsageAttribute.Set(ctx, usage); usageErr != nil { + return nil, usageErr + } generation := state.CompactionGeneration + 1 if err := contextSummaryAttribute.Set(ctx, ContextSummary{ - Generation: generation, - SummarizedThroughSequence: input, - Content: summary, + Generation: generation, SummarizedThroughSequence: input, Content: summary, }); err != nil { return nil, err } @@ -2416,10 +2490,11 @@ func (step compactContextStep) Execute(ctx dex.Context, input Sequence) (*dex.St return nil, err } activity := AgentEvent{ - Kind: EventKindCompacted, - Message: fmt.Sprintf("Compacted conversation through message %d.", input), + Kind: EventKindCompacted, + Message: condenseActivityMessage(fmt.Sprintf( + "Compacted conversation through message %d.", input, + )), } - activity.Message = condenseActivityMessage(activity.Message) if err := agentActivityStream.Write(ctx, activity); err != nil { return nil, err } @@ -2463,9 +2538,7 @@ func (step callModelStep) Execute(ctx dex.Context, _ dex.None) (*dex.StepDecisio return nil, err } progress := modelProgress{ - ctx: ctx, - activityWriteFunc: step.flow.writeActivity, - messageSequence: state.NextSequence, + ctx: ctx, activityWriteFunc: step.flow.writeActivity, messageSequence: state.NextSequence, } modelStartedAt := ctx.FirstAttemptAt().UTC() attempt := ctx.Attempt() @@ -2481,44 +2554,47 @@ func (step callModelStep) Execute(ctx dex.Context, _ dex.None) (*dex.StepDecisio } return reasoningWriter.Write(chunk) } - if activityErr := progress.writeActivity(AgentEvent{Kind: EventKindModelStarted, Message: "Calling " + string(config.Model) + ".", Attempt: &attempt}); activityErr != nil { + if activityErr := progress.writeActivity(AgentEvent{ + Kind: EventKindModelStarted, Message: "Calling " + string(config.Model) + ".", Attempt: &attempt, + }); activityErr != nil { return nil, activityErr } messages, err := step.flow.contextMessages(ctx, config, state) if err != nil { return nil, err } - reply, err := step.flow.modelClient.Complete(ctx, ModelRequest{ - Config: config, - Messages: messages, - Tools: tools, - WriteAssistant: writeAssistant, - WriteReasoning: writeReasoning, - WriteActivity: progress.writeActivity, - ForcedTool: forcedToolName, - FlowID: FlowID(ctx.FlowID()), + reply, usage, err := step.flow.modelClient.Complete(ctx, ModelRequest{ + Config: config, Messages: messages, Tools: tools, + WriteAssistant: writeAssistant, WriteReasoning: writeReasoning, + WriteActivity: progress.writeActivity, ForcedTool: forcedToolName, + FlowID: FlowID(ctx.FlowID()), CallID: modelCallID(ctx, modelCallKindResponse), }) if err != nil { - if eventErr := progress.writeActivity(AgentEvent{Kind: EventKindModelFailed, Message: "Model request failed.", Attempt: &attempt}); eventErr != nil { + if eventErr := progress.writeActivity(AgentEvent{ + Kind: EventKindModelFailed, Message: "Model request failed.", Attempt: &attempt, + }); eventErr != nil { return nil, errors.Join(err, eventErr) } return nil, err } + if usageErr := modelUsageAttribute.Set(ctx, usage); usageErr != nil { + return nil, usageErr + } if strings.TrimSpace(reply.Content) == "" && len(reply.ToolCalls) == 0 { return nil, errors.New("the model returned no content or tool calls") } - messageSequence, appendErr := step.flow.appendMessage(ctx, AgentMessage{ - Role: MessageRoleAssistant, - Content: reply.Content, - ToolCalls: reply.ToolCalls, - ProviderContextItems: reply.ProviderContextItems, - StartedAt: &modelStartedAt, + messageSequence, err := step.flow.appendMessage(ctx, AgentMessage{ + Role: MessageRoleAssistant, Content: reply.Content, + ToolCalls: reply.ToolCalls, ProviderContextItems: reply.ProviderContextItems, + StartedAt: &modelStartedAt, }) - if appendErr != nil { - return nil, appendErr + if err != nil { + return nil, err } if messageSequence != progress.messageSequence { - return nil, fmt.Errorf("assistant message sequence changed from %d to %d", progress.messageSequence, messageSequence) + return nil, fmt.Errorf( + "assistant message sequence changed from %d to %d", progress.messageSequence, messageSequence, + ) } eventMessage := "Model response completed." if len(reply.ToolCalls) > 0 { @@ -2528,7 +2604,9 @@ func (step callModelStep) Execute(ctx dex.Context, _ dex.None) (*dex.StepDecisio } eventMessage = "Model requested: " + strings.Join(names, ", ") } - if activityErr := progress.writeActivity(AgentEvent{Kind: EventKindModelCompleted, Message: eventMessage, Attempt: &attempt}); activityErr != nil { + if activityErr := progress.writeActivity(AgentEvent{ + Kind: EventKindModelCompleted, Message: eventMessage, Attempt: &attempt, + }); activityErr != nil { return nil, activityErr } state, err = agentStateAttribute.Get(ctx) @@ -2574,6 +2652,437 @@ func (step callModelStep) Execute(ctx dex.Context, _ dex.None) (*dex.StepDecisio return dex.GoTo(checkSteeredStep{flow: step.flow}, continueRouteTool), nil } +type beforeModelCallStep struct { + dex.StepDefaultsNoWaitFor[modelCallPreparation] + flow *Flow +} + +var _ dex.Step[modelCallPreparation] = beforeModelCallStep{} + +func (beforeModelCallStep) GetStepType() string { return string(stepTypeBeforeModelCall) } + +func (beforeModelCallStep) GetStepOptions() *dex.StepOptions { return modelStepOptions } + +func (step beforeModelCallStep) Execute( + ctx dex.Context, + input modelCallPreparation, +) (*dex.StepDecision, error) { + if step.flow.modelHooks == nil { + return nil, errors.New("model hooks are not configured") + } + model, err := modelForCall(ctx, input.Kind) + if err != nil { + return nil, err + } + provider, err := model.Provider() + if err != nil { + return nil, err + } + if provider != ProviderOpenAI && provider != ProviderMock { + return nil, fmt.Errorf("model hooks require an exact-usage provider, got %q", provider) + } + callID := modelCallID(ctx, input.Kind) + admitted, err := step.flow.modelHooks.BeforeModelCall( + ctx, + FlowID(ctx.FlowID()), + callID, + model, + input.Kind == modelCallKindCompaction, + ) + if err != nil { + return nil, err + } + if !admitted { + if err := step.flow.writeActivity(ctx, AgentEvent{ + Kind: EventKindModelFailed, + Message: "Model call was not admitted.", + }); err != nil { + return nil, err + } + return dex.GoTo(awaitUserStep{flow: step.flow}, nil), nil + } + return dex.GoTo(callModelWithHooksStep{flow: step.flow}, modelCallInput{ + CallID: callID, + Kind: input.Kind, + ThroughSequence: input.ThroughSequence, + }), nil +} + +type callModelWithHooksStep struct { + dex.StepDefaultsNoWaitFor[modelCallInput] + flow *Flow +} + +var _ dex.Step[modelCallInput] = callModelWithHooksStep{} + +func (callModelWithHooksStep) GetStepType() string { return string(stepTypeCallModelWithHooks) } + +func (callModelWithHooksStep) GetStepOptions() *dex.StepOptions { return modelStepOptions } + +func (step callModelWithHooksStep) Execute( + ctx dex.Context, + input modelCallInput, +) (*dex.StepDecision, error) { + var result modelCallResult + var err error + switch input.Kind { + case modelCallKindResponse: + result, err = step.flow.callResponseModel(ctx, input.CallID) + case modelCallKindCompaction: + result, err = step.flow.callCompactionModel(ctx, input.CallID, input.ThroughSequence) + default: + err = fmt.Errorf("unknown model call kind %q", input.Kind) + } + if err != nil { + return nil, err + } + return dex.GoTo(afterModelCallStep{flow: step.flow}, result), nil +} + +type afterModelCallStep struct { + dex.StepDefaultsNoWaitFor[modelCallResult] + flow *Flow +} + +var _ dex.Step[modelCallResult] = afterModelCallStep{} + +func (afterModelCallStep) GetStepType() string { return string(stepTypeAfterModelCall) } + +func (afterModelCallStep) GetStepOptions() *dex.StepOptions { return modelStepOptions } + +func (step afterModelCallStep) Execute( + ctx dex.Context, + input modelCallResult, +) (*dex.StepDecision, error) { + if step.flow.modelHooks == nil { + return nil, errors.New("model hooks are not configured") + } + usage, err := modelUsageAttribute.Get(ctx) + if err != nil { + return nil, err + } + if err := step.flow.modelHooks.AfterModelCall( + ctx, + FlowID(ctx.FlowID()), + input.CallID, + usage, + ); err != nil { + return nil, err + } + return dex.GoTo(applyModelResultStep{flow: step.flow}, input), nil +} + +type applyModelResultStep struct { + dex.StepDefaultsNoWaitFor[modelCallResult] + flow *Flow +} + +var _ dex.Step[modelCallResult] = applyModelResultStep{} + +func (applyModelResultStep) GetStepType() string { return string(stepTypeApplyModelResult) } + +func (applyModelResultStep) GetStepOptions() *dex.StepOptions { return modelStepOptions } + +func (step applyModelResultStep) Execute( + ctx dex.Context, + input modelCallResult, +) (*dex.StepDecision, error) { + switch input.Kind { + case modelCallKindResponse: + if input.Response == nil || input.Compaction != nil { + return nil, errors.New("response model result is invalid") + } + next, err := step.flow.applyResponseModelResult(ctx, *input.Response) + if err != nil { + return nil, err + } + return dex.GoTo(checkSteeredStep{flow: step.flow}, next), nil + case modelCallKindCompaction: + if input.Compaction == nil || input.Response != nil { + return nil, errors.New("compaction model result is invalid") + } + next, err := step.flow.applyCompactionModelResult(ctx, *input.Compaction) + if err != nil { + return nil, err + } + return dex.GoTo(checkSteeredStep{flow: step.flow}, next), nil + default: + return nil, fmt.Errorf("unknown model call kind %q", input.Kind) + } +} + +func modelForCall(ctx dex.Context, kind modelCallKind) (Model, error) { + config, err := agentConfigAttribute.Get(ctx) + if err != nil { + return "", err + } + switch kind { + case modelCallKindResponse: + return config.Model, nil + case modelCallKindCompaction: + if config.CompactionModel != nil { + return *config.CompactionModel, nil + } + return config.Model, nil + default: + return "", fmt.Errorf("unknown model call kind %q", kind) + } +} + +func (flow *Flow) callCompactionModel( + ctx dex.Context, + callID CallID, + throughSequence Sequence, +) (modelCallResult, error) { + if err := flow.updateStatus(ctx, AgentStatusCompactingContext); err != nil { + return modelCallResult{}, err + } + config, err := agentConfigAttribute.Get(ctx) + if err != nil { + return modelCallResult{}, err + } + state, err := agentStateAttribute.Get(ctx) + if err != nil { + return modelCallResult{}, err + } + messages, err := flow.loadMessages(ctx, state.SummarizedThroughSequence+1, throughSequence, config) + if err != nil { + return modelCallResult{}, err + } + previousSummary, err := flow.getSummary(ctx) + if err != nil { + return modelCallResult{}, err + } + if heartbeatErr := ctx.RecordHeartbeat(compactionHeartbeat{ + Phase: heartbeatPhaseCompacting, + ThroughSequence: throughSequence, + }); heartbeatErr != nil { + return modelCallResult{}, heartbeatErr + } + summary, usage, err := flow.modelClient.Summarize(ctx, SummarizeRequest{ + Config: config, + PreviousSummary: previousSummary.Content, + Messages: messages, + FlowID: FlowID(ctx.FlowID()), + CallID: callID, + }) + if err != nil { + if eventErr := flow.writeActivity(ctx, AgentEvent{ + Kind: EventKindCompactionFailed, Message: "Context compaction failed.", + }); eventErr != nil { + return modelCallResult{}, errors.Join(err, eventErr) + } + return modelCallResult{}, err + } + if usageErr := modelUsageAttribute.Set(ctx, usage); usageErr != nil { + return modelCallResult{}, usageErr + } + return modelCallResult{ + CallID: callID, + Kind: modelCallKindCompaction, + Compaction: &modelCompactionResult{ + Summary: summary, ThroughSequence: throughSequence, + }, + }, nil +} + +func (flow *Flow) callResponseModel(ctx dex.Context, callID CallID) (modelCallResult, error) { + if err := flow.updateStatus(ctx, AgentStatusCallingModel); err != nil { + return modelCallResult{}, err + } + config, err := agentConfigAttribute.Get(ctx) + if err != nil { + return modelCallResult{}, err + } + state, err := agentStateAttribute.Get(ctx) + if err != nil { + return modelCallResult{}, err + } + tools := flow.invocationToolDefinitions(config, state) + var forcedToolName ToolName + if state.InteractionMode == InteractionModePlanning && state.PlanningRequiresWrite { + forcedToolName = ToolNameWriteTodos + } + assistantWriter, err := dex.NewBufferedTextStream(ctx, assistantTextStream) + if err != nil { + return modelCallResult{}, err + } + reasoningWriter, err := dex.NewBufferedTextStream(ctx, reasoningSummaryStream) + if err != nil { + return modelCallResult{}, err + } + progress := modelProgress{ + ctx: ctx, activityWriteFunc: flow.writeActivity, messageSequence: state.NextSequence, + } + modelStartedAt := ctx.FirstAttemptAt().UTC() + attempt := ctx.Attempt() + writeAssistant := func(chunk string) error { + if heartbeatErr := ctx.RecordHeartbeat(modelHeartbeat{Phase: heartbeatPhaseAssistantStream}); heartbeatErr != nil { + return heartbeatErr + } + return assistantWriter.Write(chunk) + } + writeReasoning := func(chunk string) error { + if heartbeatErr := ctx.RecordHeartbeat(modelHeartbeat{Phase: heartbeatPhaseReasoningStream}); heartbeatErr != nil { + return heartbeatErr + } + return reasoningWriter.Write(chunk) + } + if activityErr := progress.writeActivity(AgentEvent{ + Kind: EventKindModelStarted, Message: "Calling " + string(config.Model) + ".", Attempt: &attempt, + }); activityErr != nil { + return modelCallResult{}, activityErr + } + messages, err := flow.contextMessages(ctx, config, state) + if err != nil { + return modelCallResult{}, err + } + reply, usage, err := flow.modelClient.Complete(ctx, ModelRequest{ + Config: config, Messages: messages, Tools: tools, + WriteAssistant: writeAssistant, WriteReasoning: writeReasoning, + WriteActivity: progress.writeActivity, ForcedTool: forcedToolName, + FlowID: FlowID(ctx.FlowID()), CallID: callID, + }) + if err != nil { + if eventErr := progress.writeActivity(AgentEvent{ + Kind: EventKindModelFailed, Message: "Model request failed.", Attempt: &attempt, + }); eventErr != nil { + return modelCallResult{}, errors.Join(err, eventErr) + } + return modelCallResult{}, err + } + if strings.TrimSpace(reply.Content) == "" && len(reply.ToolCalls) == 0 { + return modelCallResult{}, errors.New("the model returned no content or tool calls") + } + if usageErr := modelUsageAttribute.Set(ctx, usage); usageErr != nil { + return modelCallResult{}, usageErr + } + eventMessage := "Model response completed." + if len(reply.ToolCalls) > 0 { + names := make([]string, 0, len(reply.ToolCalls)) + for _, call := range reply.ToolCalls { + names = append(names, string(call.Name)) + } + eventMessage = "Model requested: " + strings.Join(names, ", ") + } + if err := progress.writeActivity(AgentEvent{ + Kind: EventKindModelCompleted, Message: eventMessage, Attempt: &attempt, + }); err != nil { + return modelCallResult{}, err + } + return modelCallResult{ + CallID: callID, + Kind: modelCallKindResponse, + Response: &modelResponseResult{ + Reply: reply, StartedAt: modelStartedAt, ExpectedMessageSequence: progress.messageSequence, + }, + }, nil +} + +func (flow *Flow) applyCompactionModelResult( + ctx dex.Context, + result modelCompactionResult, +) (continuation, error) { + config, err := agentConfigAttribute.Get(ctx) + if err != nil { + return "", err + } + state, err := agentStateAttribute.Get(ctx) + if err != nil { + return "", err + } + generation := state.CompactionGeneration + 1 + if err := contextSummaryAttribute.Set(ctx, ContextSummary{ + Generation: generation, SummarizedThroughSequence: result.ThroughSequence, Content: result.Summary, + }); err != nil { + return "", err + } + state.SummarizedThroughSequence = result.ThroughSequence + state.CompactionGeneration = generation + if err := agentStateAttribute.Set(ctx, state); err != nil { + return "", err + } + if _, err := flow.trimSummarizedMessages(ctx, config, state); err != nil { + return "", err + } + activity := AgentEvent{ + Kind: EventKindCompacted, + Message: condenseActivityMessage(fmt.Sprintf( + "Compacted conversation through message %d.", result.ThroughSequence, + )), + } + if err := agentActivityStream.Write(ctx, activity); err != nil { + return "", err + } + return continueCallModel, nil +} + +func (flow *Flow) applyResponseModelResult( + ctx dex.Context, + result modelResponseResult, +) (continuation, error) { + messageSequence, err := flow.appendMessage(ctx, AgentMessage{ + Role: MessageRoleAssistant, Content: result.Reply.Content, + ToolCalls: result.Reply.ToolCalls, ProviderContextItems: result.Reply.ProviderContextItems, + StartedAt: &result.StartedAt, + }) + if err != nil { + return "", err + } + if messageSequence != result.ExpectedMessageSequence { + return "", fmt.Errorf( + "assistant message sequence changed from %d to %d", + result.ExpectedMessageSequence, + messageSequence, + ) + } + config, err := agentConfigAttribute.Get(ctx) + if err != nil { + return "", err + } + state, err := agentStateAttribute.Get(ctx) + if err != nil { + return "", err + } + state, err = flow.trimSummarizedMessages(ctx, config, state) + if err != nil { + return "", err + } + if len(result.Reply.ToolCalls) == 0 { + plan, err := flow.getPlan(ctx) + if err != nil { + return "", err + } + if plan == nil || plan.Status == PlanStatusCompleted { + state.InteractionMode = InteractionModeChat + state.PlanningRequiresWrite = false + state.PlanNoProgressAttempts = 0 + if err := agentStateAttribute.Set(ctx, state); err != nil { + return "", err + } + } else if plan.Status == PlanStatusActive && + state.InteractionMode == InteractionModeExecuting && + !allTasksCompleted(plan.Tasks) { + state.PlanNoProgressAttempts++ + if err := agentStateAttribute.Set(ctx, state); err != nil { + return "", err + } + if state.PlanNoProgressAttempts <= maximumAutomaticPlanRecoveryCount { + return continueCallModel, nil + } + } + return continueAwaitUser, nil + } + state.Status = AgentStatusRoutingTool + state.PendingToolCalls = result.Reply.ToolCalls + state.PendingToolIndex = 0 + state.PlanNoProgressAttempts = 0 + if err := agentStateAttribute.Set(ctx, state); err != nil { + return "", err + } + return continueRouteTool, nil +} + type checkSteeredStep struct { dex.StepDefaults flow *Flow @@ -2607,13 +3116,25 @@ func (step checkSteeredStep) Execute(ctx dex.Context, input continuation) (*dex. return nil, err } if cutoff != nil { - return dex.GoTo(compactContextStep{flow: step.flow}, *cutoff), nil + if step.flow.modelHooks == nil { + return dex.GoTo(compactContextStep{flow: step.flow}, *cutoff), nil + } + return dex.GoTo(beforeModelCallStep{flow: step.flow}, modelCallPreparation{ + Kind: modelCallKindCompaction, + ThroughSequence: *cutoff, + }), nil } - return dex.GoTo(callModelStep{flow: step.flow}, nil), nil + if step.flow.modelHooks == nil { + return dex.GoTo(callModelStep{flow: step.flow}, nil), nil + } + return dex.GoTo(beforeModelCallStep{flow: step.flow}, modelCallPreparation{Kind: modelCallKindResponse}), nil case continueAwaitUser: return dex.GoTo(awaitUserStep{flow: step.flow}, nil), nil case continueCallModel: - return dex.GoTo(callModelStep{flow: step.flow}, nil), nil + if step.flow.modelHooks == nil { + return dex.GoTo(callModelStep{flow: step.flow}, nil), nil + } + return dex.GoTo(beforeModelCallStep{flow: step.flow}, modelCallPreparation{Kind: modelCallKindResponse}), nil case continueRouteTool: return dex.GoTo(routeToolStep{flow: step.flow}, nil), nil case continueAwaitToolApproval: diff --git a/internal/agent/flow_integration_test.go b/internal/agent/flow_integration_test.go index 6127123..6b93e4c 100644 --- a/internal/agent/flow_integration_test.go +++ b/internal/agent/flow_integration_test.go @@ -731,6 +731,243 @@ func TestAgentRejectsInvalidWriteTodosIntegration(t *testing.T) { } } +func TestAgentModelHooksIntegration(t *testing.T) { + t.Run("retry preserves CallID and usage is not accumulated", func(t *testing.T) { + modelClient := &usageIntegrationModel{} + var mutex sync.Mutex + beforeCallIDs := make([]CallID, 0, 3) + afterCallIDs := make([]CallID, 0, 2) + afterUsages := make([]ModelUsage, 0, 2) + beforeAttempts := 0 + afterAttempts := 0 + hooks := &ModelHooksConfig{ + BeforeModelCall: func(_ context.Context, _ FlowID, callID CallID, _ Model, isCompaction bool) (bool, error) { + if isCompaction { + return false, errors.New("unexpected compaction") + } + mutex.Lock() + defer mutex.Unlock() + beforeCallIDs = append(beforeCallIDs, callID) + beforeAttempts++ + if beforeAttempts == 1 { + return false, errors.New("retry admission") + } + return true, nil + }, + AfterModelCall: func(_ context.Context, _ FlowID, callID CallID, usage ModelUsage) error { + mutex.Lock() + defer mutex.Unlock() + afterCallIDs = append(afterCallIDs, callID) + afterUsages = append(afterUsages, usage) + afterAttempts++ + if afterAttempts == 1 { + return errors.New("retry usage synchronization") + } + return nil + }, + } + environment := newAgentIntegrationEnvironment( + t, modelClient, newIntegrationToolRegistry(), WithModelHooks(hooks), + ) + flowID := FlowID("agent-model-hooks-" + randomLocalID(t)) + if _, err := environment.agent.Start(t.Context(), flowID, StartRequest{Config: NewAgentConfig()}); err != nil { + t.Fatal(err) + } + waitForAgentState(t, environment, flowID, func(state AgentState) bool { + return state.Status == AgentStatusWaitingForMessage + }) + for _, message := range []string{"first", "second"} { + if err := environment.agent.SendMessage(t.Context(), flowID, UserMessage{Content: message}); err != nil { + t.Fatal(err) + } + waitForSnapshot(t, environment, flowID, func(snapshot AgentSnapshot) bool { + return snapshot.Description.Status == AgentStatusWaitingForMessage && + historyHasMessage(snapshot.History.Messages, MessageRoleAssistant, "integration response: "+message) + }) + } + mutex.Lock() + defer mutex.Unlock() + if len(beforeCallIDs) != 3 || beforeCallIDs[0] == "" || beforeCallIDs[0] != beforeCallIDs[1] || beforeCallIDs[2] == beforeCallIDs[1] { + t.Fatalf("before CallIDs = %q", beforeCallIDs) + } + if len(afterCallIDs) != 3 || afterCallIDs[0] != beforeCallIDs[1] || + afterCallIDs[1] != beforeCallIDs[1] || afterCallIDs[2] != beforeCallIDs[2] { + t.Fatalf("after CallIDs = %q, before = %q", afterCallIDs, beforeCallIDs) + } + if len(afterUsages) != 3 || afterUsages[0].TotalTokens != 1 || + afterUsages[1].TotalTokens != 1 || afterUsages[2].TotalTokens != 2 { + t.Fatalf("after usage = %#v", afterUsages) + } + if modelClient.completeCalls.Load() != 2 { + t.Fatalf("model calls = %d, want 2", modelClient.completeCalls.Load()) + } + }) + + t.Run("Worker replacement after Call retries only After", func(t *testing.T) { + modelClient := &usageIntegrationModel{} + afterStarted := make(chan struct{}) + var afterCalls atomic.Int32 + hooks := &ModelHooksConfig{ + BeforeModelCall: func(context.Context, FlowID, CallID, Model, bool) (bool, error) { + return true, nil + }, + AfterModelCall: func(_ context.Context, _ FlowID, _ CallID, _ ModelUsage) error { + if afterCalls.Add(1) == 1 { + close(afterStarted) + return errors.New("replace Worker before After retry") + } + return nil + }, + } + environment := newAgentIntegrationEnvironment( + t, modelClient, newIntegrationToolRegistry(), WithModelHooks(hooks), + ) + flowID := FlowID("agent-model-hooks-replace-" + randomLocalID(t)) + if _, err := environment.agent.Start(t.Context(), flowID, StartRequest{Config: NewAgentConfig()}); err != nil { + t.Fatal(err) + } + waitForAgentState(t, environment, flowID, func(state AgentState) bool { + return state.Status == AgentStatusWaitingForMessage + }) + if err := environment.agent.SendMessage(t.Context(), flowID, UserMessage{Content: "replace"}); err != nil { + t.Fatal(err) + } + select { + case <-afterStarted: + case <-time.After(integrationWaitTimeout): + t.Fatal("After hook did not start") + } + environment.replaceWorker(t, flowID) + waitForSnapshot(t, environment, flowID, func(snapshot AgentSnapshot) bool { + return snapshot.Description.Status == AgentStatusWaitingForMessage && + historyHasMessage(snapshot.History.Messages, MessageRoleAssistant, "integration response: replace") + }) + if modelClient.completeCalls.Load() != 1 || afterCalls.Load() < 2 { + t.Fatalf("model/after calls = %d/%d, want 1/at least 2", modelClient.completeCalls.Load(), afterCalls.Load()) + } + }) + + t.Run("compaction uses hooks", func(t *testing.T) { + modelClient := &usageIntegrationModel{} + var compactionBefore atomic.Int32 + hooks := &ModelHooksConfig{ + BeforeModelCall: func(_ context.Context, _ FlowID, _ CallID, _ Model, isCompaction bool) (bool, error) { + if isCompaction { + compactionBefore.Add(1) + } + return true, nil + }, + AfterModelCall: func(context.Context, FlowID, CallID, ModelUsage) error { return nil }, + } + environment := newAgentIntegrationEnvironment( + t, modelClient, newIntegrationToolRegistry(), WithModelHooks(hooks), + ) + flowID := FlowID("agent-model-hooks-compaction-" + randomLocalID(t)) + config := NewAgentConfig() + config.MaxContextTokens = 80 + config.CompactionTriggerFraction = 0.60 + config.CompactionKeepFraction = 0.20 + if _, err := environment.agent.Start(t.Context(), flowID, StartRequest{Config: config}); err != nil { + t.Fatal(err) + } + waitForAgentState(t, environment, flowID, func(state AgentState) bool { + return state.Status == AgentStatusWaitingForMessage + }) + for index := 0; index < 6 && compactionBefore.Load() == 0; index++ { + message := fmt.Sprintf("long-%d-%s", index, strings.Repeat("context ", 8)) + previousSequence := readSnapshot(t, environment, flowID).Description.LastSequence + previousCompactions := compactionBefore.Load() + if err := environment.agent.SendMessage(t.Context(), flowID, UserMessage{Content: message}); err != nil { + t.Fatal(err) + } + state := waitForAgentState(t, environment, flowID, func(state AgentState) bool { + return state.Status == AgentStatusWaitingForMessage && state.LastSequence >= previousSequence+2 + }) + wantResponse := "integration response: " + message + if compactionBefore.Load() > previousCompactions { + wantResponse = "integration response: " + } + assertApplicationMessage(t, environment, flowID, state.LastSequence, MessageRoleAssistant, wantResponse) + } + if compactionBefore.Load() == 0 || modelClient.summarizeCalls.Load() == 0 { + t.Fatalf("compaction before/model calls = %d/%d", compactionBefore.Load(), modelClient.summarizeCalls.Load()) + } + }) + + t.Run("business rejection skips model and after hook", func(t *testing.T) { + modelClient := &usageIntegrationModel{} + var afterCalls atomic.Int32 + hooks := &ModelHooksConfig{ + BeforeModelCall: func(context.Context, FlowID, CallID, Model, bool) (bool, error) { + return false, nil + }, + AfterModelCall: func(context.Context, FlowID, CallID, ModelUsage) error { + afterCalls.Add(1) + return nil + }, + } + environment := newAgentIntegrationEnvironment( + t, modelClient, newIntegrationToolRegistry(), WithModelHooks(hooks), + ) + flowID := FlowID("agent-model-hooks-rejected-" + randomLocalID(t)) + if _, err := environment.agent.Start(t.Context(), flowID, StartRequest{Config: NewAgentConfig()}); err != nil { + t.Fatal(err) + } + waitForAgentState(t, environment, flowID, func(state AgentState) bool { + return state.Status == AgentStatusWaitingForMessage + }) + if err := environment.agent.SendMessage(t.Context(), flowID, UserMessage{Content: "rejected"}); err != nil { + t.Fatal(err) + } + waitForSnapshot(t, environment, flowID, func(snapshot AgentSnapshot) bool { + return snapshot.Description.Status == AgentStatusWaitingForMessage && + historyHasMessage(snapshot.History.Messages, MessageRoleUser, "rejected") + }) + if modelClient.completeCalls.Load() != 0 || afterCalls.Load() != 0 { + t.Fatalf("model/after calls = %d/%d, want 0/0", modelClient.completeCalls.Load(), afterCalls.Load()) + } + }) + + t.Run("provider without exact usage fails before dispatch", func(t *testing.T) { + modelClient := &usageIntegrationModel{} + var beforeCalls atomic.Int32 + hooks := &ModelHooksConfig{ + BeforeModelCall: func(context.Context, FlowID, CallID, Model, bool) (bool, error) { + beforeCalls.Add(1) + return true, nil + }, + AfterModelCall: func(context.Context, FlowID, CallID, ModelUsage) error { return nil }, + } + environment := newAgentIntegrationEnvironment( + t, modelClient, newIntegrationToolRegistry(), WithModelHooks(hooks), + ) + flowID := FlowID("agent-model-hooks-provider-" + randomLocalID(t)) + config := NewAgentConfig() + config.Model = "anthropic/claude-test" + if _, err := environment.agent.Start(t.Context(), flowID, StartRequest{Config: config}); err != nil { + t.Fatal(err) + } + waitForAgentState(t, environment, flowID, func(state AgentState) bool { + return state.Status == AgentStatusWaitingForMessage + }) + if err := environment.agent.SendMessage(t.Context(), flowID, UserMessage{Content: "unsupported"}); err != nil { + t.Fatal(err) + } + ctx, cancel := context.WithTimeout(t.Context(), integrationWaitTimeout) + defer cancel() + result, err := environment.sdk.WaitForFlow(ctx, string(flowID), dex.WaitForFlowOptions{}) + if err != nil { + t.Fatal(err) + } + if result.Status != dex.FlowFailed { + t.Fatalf("Flow status = %v, want failed", result.Status) + } + if beforeCalls.Load() != 0 || modelClient.completeCalls.Load() != 0 { + t.Fatalf("before/model calls = %d/%d, want 0/0", beforeCalls.Load(), modelClient.completeCalls.Load()) + } + }) +} + func TestAgentFlowDurabilityIntegration(t *testing.T) { modelClient := integrationModel{} toolRegistry := newIntegrationToolRegistry() @@ -2520,7 +2757,43 @@ type integrationModel struct{} var _ ModelClient = integrationModel{} -func (integrationModel) Complete(ctx context.Context, request ModelRequest) (ModelReply, error) { +type usageIntegrationModel struct { + integrationModel + completeCalls atomic.Int32 + summarizeCalls atomic.Int32 +} + +func (model *usageIntegrationModel) Complete( + ctx context.Context, + request ModelRequest, +) (ModelReply, ModelUsage, error) { + reply, _, err := model.integrationModel.Complete(ctx, request) + if err != nil { + return ModelReply{}, ModelUsage{}, err + } + call := int64(model.completeCalls.Add(1)) + return reply, ModelUsage{InputTokens: call, TotalTokens: call}, nil +} + +func (model *usageIntegrationModel) Summarize( + ctx context.Context, + request SummarizeRequest, +) (string, ModelUsage, error) { + summary, _, err := model.integrationModel.Summarize(ctx, request) + if err != nil { + return "", ModelUsage{}, err + } + call := int64(model.summarizeCalls.Add(1)) + return summary, ModelUsage{InputTokens: call, TotalTokens: call}, nil +} + +func (model integrationModel) Complete(ctx context.Context, request ModelRequest) (ModelReply, ModelUsage, error) { + reply, err := model.complete(ctx, request) + usage := ModelUsage{InputTokens: 11, OutputTokens: 7, TotalTokens: 18} + return reply, usage, err +} + +func (integrationModel) complete(ctx context.Context, request ModelRequest) (ModelReply, error) { if request.WriteAssistant == nil || request.WriteReasoning == nil || request.WriteActivity == nil { return ModelReply{}, errors.New("integration model writers are required") } @@ -2662,7 +2935,12 @@ func (integrationModel) Complete(ctx context.Context, request ModelRequest) (Mod return ModelReply{Content: content, ToolCalls: []ToolCall{}}, nil } -func (integrationModel) Summarize(_ context.Context, request SummarizeRequest) (string, error) { +func (model integrationModel) Summarize(ctx context.Context, request SummarizeRequest) (string, ModelUsage, error) { + summary, err := model.summarize(ctx, request) + return summary, ModelUsage{InputTokens: 13, OutputTokens: 5, TotalTokens: 18}, err +} + +func (integrationModel) summarize(_ context.Context, request SummarizeRequest) (string, error) { parts := make([]string, 0, len(request.Messages)+1) if request.PreviousSummary != "" { parts = append(parts, request.PreviousSummary) @@ -2701,7 +2979,7 @@ func newBlockingFirstModel() *blockingFirstModel { return &blockingFirstModel{started: make(chan struct{}), release: make(chan struct{})} } -func (model *blockingFirstModel) Complete(ctx context.Context, request ModelRequest) (ModelReply, error) { +func (model *blockingFirstModel) Complete(ctx context.Context, request ModelRequest) (ModelReply, ModelUsage, error) { shouldBlock := false model.once.Do(func() { shouldBlock = true @@ -2711,7 +2989,7 @@ func (model *blockingFirstModel) Complete(ctx context.Context, request ModelRequ select { case <-model.release: case <-ctx.Done(): - return ModelReply{}, ctx.Err() + return ModelReply{}, ModelUsage{}, ctx.Err() } } return model.integrationModel.Complete(ctx, request) @@ -2726,7 +3004,7 @@ func newBlockingPlanModel() *blockingPlanModel { } } -func (model *blockingPlanModel) Complete(ctx context.Context, request ModelRequest) (ModelReply, error) { +func (model *blockingPlanModel) Complete(ctx context.Context, request ModelRequest) (ModelReply, ModelUsage, error) { if _, hasActivePlan := integrationActivePlanTaskStatus(request.Messages); !hasActivePlan { return model.integrationModel.Complete(ctx, request) } @@ -2736,14 +3014,14 @@ func (model *blockingPlanModel) Complete(ctx context.Context, request ModelReque select { case <-model.release: case <-ctx.Done(): - return ModelReply{}, ctx.Err() + return ModelReply{}, ModelUsage{}, ctx.Err() } } content := "integration stopped before completing the active plan" if err := request.WriteAssistant(content); err != nil { - return ModelReply{}, err + return ModelReply{}, ModelUsage{}, err } - return ModelReply{Content: content, ToolCalls: []ToolCall{}}, nil + return ModelReply{Content: content, ToolCalls: []ToolCall{}}, ModelUsage{InputTokens: 11, OutputTokens: 7, TotalTokens: 18}, nil } func integrationToolReply(request ModelRequest, name ToolName, arguments JSONObject, content string) (ModelReply, error) { @@ -2871,7 +3149,7 @@ type parallelIntegrationModel struct { integrationModel } -func (parallelIntegrationModel) Complete(ctx context.Context, request ModelRequest) (ModelReply, error) { +func (parallelIntegrationModel) Complete(ctx context.Context, request ModelRequest) (ModelReply, ModelUsage, error) { if integrationLastUserContent(request.Messages) != "/parallel" { return integrationModel{}.Complete(ctx, request) } @@ -2879,7 +3157,7 @@ func (parallelIntegrationModel) Complete(ctx context.Context, request ModelReque return integrationModel{}.Complete(ctx, request) } if err := request.WriteAssistant("reading in parallel"); err != nil { - return ModelReply{}, err + return ModelReply{}, ModelUsage{}, err } return ModelReply{ Content: "reading in parallel", @@ -2887,14 +3165,14 @@ func (parallelIntegrationModel) Complete(ctx context.Context, request ModelReque integrationToolCall(request, "parallel_read_a", MustJSONObject(`{}`)), integrationToolCall(request, "parallel_read_b", MustJSONObject(`{}`)), }, - }, nil + }, ModelUsage{InputTokens: 11, OutputTokens: 7, TotalTokens: 18}, nil } type mixedParallelIntegrationModel struct { integrationModel } -func (mixedParallelIntegrationModel) Complete(ctx context.Context, request ModelRequest) (ModelReply, error) { +func (mixedParallelIntegrationModel) Complete(ctx context.Context, request ModelRequest) (ModelReply, ModelUsage, error) { if integrationLastUserContent(request.Messages) != "/parallel-recovery" { return integrationModel{}.Complete(ctx, request) } @@ -2902,7 +3180,7 @@ func (mixedParallelIntegrationModel) Complete(ctx context.Context, request Model return integrationModel{}.Complete(ctx, request) } if err := request.WriteAssistant("reading mixed parallel results"); err != nil { - return ModelReply{}, err + return ModelReply{}, ModelUsage{}, err } names := []ToolName{ "parallel_success", @@ -2914,7 +3192,7 @@ func (mixedParallelIntegrationModel) Complete(ctx context.Context, request Model for _, name := range names { calls = append(calls, integrationToolCall(request, name, MustJSONObject(`{}`))) } - return ModelReply{Content: "reading mixed parallel results", ToolCalls: calls}, nil + return ModelReply{Content: "reading mixed parallel results", ToolCalls: calls}, ModelUsage{InputTokens: 11, OutputTokens: 7, TotalTokens: 18}, nil } type mixedParallelToolRegistryForTestOnly struct{} diff --git a/internal/agent/model_hooks_test.go b/internal/agent/model_hooks_test.go new file mode 100644 index 0000000..d72645f --- /dev/null +++ b/internal/agent/model_hooks_test.go @@ -0,0 +1,73 @@ +// Copyright (c) 2022-2026 Super Durable, Inc. +// +// Licensed under the Apache License, Version 2.0 (the "License"); +// you may not use this file except in compliance with the License. +// You may obtain a copy of the License at +// +// http://www.apache.org/licenses/LICENSE-2.0 +// +// Unless required by applicable law or agreed to in writing, software +// distributed under the License is distributed on an "AS IS" BASIS, +// WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. +// See the License for the specific language governing permissions and +// limitations under the License. +// +// SPDX-License-Identifier: Apache-2.0 + +package agent + +import ( + "context" + "testing" +) + +func TestNewFlowRejectsIncompleteModelHooks(t *testing.T) { + t.Parallel() + for _, config := range []*ModelHooksConfig{ + {BeforeModelCall: func(context.Context, FlowID, CallID, Model, bool) (bool, error) { return true, nil }}, + {AfterModelCall: func(context.Context, FlowID, CallID, ModelUsage) error { return nil }}, + } { + config := config + t.Run("incomplete", func(t *testing.T) { + t.Parallel() + defer func() { + if recover() == nil { + t.Fatal("NewFlow panic = nil") + } + }() + NewFlow(hookTestModel{}, hookTestTools{}, WithModelHooks(config)) + }) + } +} + +func TestNewFlowAcceptsNilModelHooks(t *testing.T) { + t.Parallel() + flow := NewFlow(hookTestModel{}, hookTestTools{}, WithModelHooks(nil)) + if flow.modelHooks != nil { + t.Fatalf("model hooks = %#v, want nil", flow.modelHooks) + } +} + +type hookTestModel struct{} + +func (hookTestModel) Complete(context.Context, ModelRequest) (ModelReply, ModelUsage, error) { + return ModelReply{}, ModelUsage{}, nil +} + +func (hookTestModel) Summarize(context.Context, SummarizeRequest) (string, ModelUsage, error) { + return "", ModelUsage{}, nil +} + +func (hookTestModel) CountTokens(Model, []AgentMessage) int { return 0 } + +type hookTestTools struct{} + +func (hookTestTools) ServerNames() []string { return nil } + +func (hookTestTools) RegisteredTools() []RegisteredTool { return nil } + +func (hookTestTools) Definitions([]string, []ToolName) []ToolDefinition { return nil } + +func (hookTestTools) Execute(context.Context, ToolInvocation) (ToolExecutionResult, error) { + return ToolExecutionResult{}, nil +} diff --git a/internal/agent/types.go b/internal/agent/types.go index 991ebee..9c92ced 100644 --- a/internal/agent/types.go +++ b/internal/agent/types.go @@ -1143,6 +1143,15 @@ type ModelReply struct { ProviderContextItems []ProviderContextItem `json:"provider_context_items"` } +// ModelUsage records the exact token usage reported by one model call. +type ModelUsage struct { + InputTokens int64 `json:"input_tokens"` + CachedInputTokens int64 `json:"cached_input_tokens"` + OutputTokens int64 `json:"output_tokens"` + ReasoningTokens int64 `json:"reasoning_tokens"` + TotalTokens int64 `json:"total_tokens"` +} + // ToolDefinition is one provider-neutral function schema and execution policy. type ToolDefinition struct { Name ToolName @@ -1180,6 +1189,7 @@ type ModelRequest struct { WriteActivity ActivityWriter ForcedTool ToolName FlowID FlowID + CallID CallID } // SummarizeRequest contains one cumulative compaction invocation. @@ -1188,15 +1198,28 @@ type SummarizeRequest struct { PreviousSummary string Messages []AgentMessage FlowID FlowID + CallID CallID } // ModelClient is the provider-neutral model boundary. type ModelClient interface { - Complete(context.Context, ModelRequest) (ModelReply, error) - Summarize(context.Context, SummarizeRequest) (string, error) + Complete(context.Context, ModelRequest) (ModelReply, ModelUsage, error) + Summarize(context.Context, SummarizeRequest) (string, ModelUsage, error) CountTokens(Model, []AgentMessage) int } +// BeforeModelCallHook admits or rejects one model call before provider dispatch. +type BeforeModelCallHook func(context.Context, FlowID, CallID, Model, bool) (bool, error) + +// AfterModelCallHook synchronizes exact usage after one successful model call. +type AfterModelCallHook func(context.Context, FlowID, CallID, ModelUsage) error + +// ModelHooksConfig configures optional durable boundaries around model calls. +type ModelHooksConfig struct { + BeforeModelCall BeforeModelCallHook + AfterModelCall AfterModelCallHook +} + // ToolInvocation contains one trusted-registry execution request. type ToolInvocation struct { FlowID FlowID diff --git a/internal/app/app.go b/internal/app/app.go index 466c3ba..d995d7d 100644 --- a/internal/app/app.go +++ b/internal/app/app.go @@ -306,6 +306,7 @@ func buildModelClient(section *config.Providers, credentials *model.CredentialSt model.NewMockClient(), model.NewOpenAIClient(credentials, httpClient, section.OpenAI.BaseURL, &model.OpenAIClientConfig{ ReasoningSummaryStreamingEnabled: true, + DisableModelResponseStore: section.OpenAIDisableModelResponseStore, }), anthropic, gemini, diff --git a/internal/config/config.go b/internal/config/config.go index 75df1b3..078d44f 100644 --- a/internal/config/config.go +++ b/internal/config/config.go @@ -32,23 +32,24 @@ import ( type EnvironmentVariable string const ( - EnvHTTPAddress EnvironmentVariable = "SUPERAGENT_HTTP_ADDRESS" - EnvHTTPAllowedOrigins EnvironmentVariable = "SUPERAGENT_HTTP_ALLOWED_ORIGINS" - EnvStreamRecoveryLimit EnvironmentVariable = "SUPERAGENT_STREAM_RECOVERY_LIMIT" - EnvDexFlowServiceAddress EnvironmentVariable = "DEX_FLOW_SERVICE_ADDRESS" - EnvDexWorkerBindAddress EnvironmentVariable = "DEX_WORKER_BIND_ADDRESS" - EnvDexWorkerTarget EnvironmentVariable = "DEX_WORKER_TARGET" - EnvBlobCacheDirectory EnvironmentVariable = "DEX_BLOB_CACHE_DIR" - EnvBlobCacheMaxBytes EnvironmentVariable = "DEX_BLOB_CACHE_MAX_BYTES" - EnvMCPConfig EnvironmentVariable = "DEX_AGENT_MCP_CONFIG" - EnvOpenAIAPIKey EnvironmentVariable = "OPENAI_API_KEY" //nolint:gosec // Environment variable name, not a credential. - EnvAnthropicAPIKey EnvironmentVariable = "ANTHROPIC_API_KEY" //nolint:gosec // Environment variable name, not a credential. - EnvGeminiAPIKey EnvironmentVariable = "GEMINI_API_KEY" //nolint:gosec // Environment variable name, not a credential. - EnvGroqAPIKey EnvironmentVariable = "GROQ_API_KEY" //nolint:gosec // Environment variable name, not a credential. - EnvOpenAIBaseURL EnvironmentVariable = "OPENAI_BASE_URL" - EnvAnthropicBaseURL EnvironmentVariable = "ANTHROPIC_BASE_URL" - EnvGeminiBaseURL EnvironmentVariable = "GEMINI_BASE_URL" - EnvGroqBaseURL EnvironmentVariable = "GROQ_BASE_URL" + EnvHTTPAddress EnvironmentVariable = "SUPERAGENT_HTTP_ADDRESS" + EnvHTTPAllowedOrigins EnvironmentVariable = "SUPERAGENT_HTTP_ALLOWED_ORIGINS" + EnvStreamRecoveryLimit EnvironmentVariable = "SUPERAGENT_STREAM_RECOVERY_LIMIT" + EnvDexFlowServiceAddress EnvironmentVariable = "DEX_FLOW_SERVICE_ADDRESS" + EnvDexWorkerBindAddress EnvironmentVariable = "DEX_WORKER_BIND_ADDRESS" + EnvDexWorkerTarget EnvironmentVariable = "DEX_WORKER_TARGET" + EnvBlobCacheDirectory EnvironmentVariable = "DEX_BLOB_CACHE_DIR" + EnvBlobCacheMaxBytes EnvironmentVariable = "DEX_BLOB_CACHE_MAX_BYTES" + EnvMCPConfig EnvironmentVariable = "DEX_AGENT_MCP_CONFIG" + EnvOpenAIAPIKey EnvironmentVariable = "OPENAI_API_KEY" //nolint:gosec // Environment variable name, not a credential. + EnvAnthropicAPIKey EnvironmentVariable = "ANTHROPIC_API_KEY" //nolint:gosec // Environment variable name, not a credential. + EnvGeminiAPIKey EnvironmentVariable = "GEMINI_API_KEY" //nolint:gosec // Environment variable name, not a credential. + EnvGroqAPIKey EnvironmentVariable = "GROQ_API_KEY" //nolint:gosec // Environment variable name, not a credential. + EnvOpenAIBaseURL EnvironmentVariable = "OPENAI_BASE_URL" + EnvOpenAIDisableModelResponseStore EnvironmentVariable = "OPENAI_DISABLE_MODEL_RESPONSE_STORE" + EnvAnthropicBaseURL EnvironmentVariable = "ANTHROPIC_BASE_URL" + EnvGeminiBaseURL EnvironmentVariable = "GEMINI_BASE_URL" + EnvGroqBaseURL EnvironmentVariable = "GROQ_BASE_URL" ) const ( @@ -114,11 +115,12 @@ type MCP struct { // Providers configures credentials and trusted API origins. type Providers struct { - OpenAI *Provider - Anthropic *Provider - Gemini *Provider - Groq *Provider - RequestTimeout time.Duration + OpenAI *Provider + Anthropic *Provider + Gemini *Provider + Groq *Provider + OpenAIDisableModelResponseStore bool + RequestTimeout time.Duration } // Provider contains one provider's process credential and optional origin override. @@ -151,6 +153,10 @@ func load(lookup Lookup) (*Config, error) { if err != nil { return nil, err } + disableOpenAIResponseStore, err := optionalBool(lookup, EnvOpenAIDisableModelResponseStore, false) + if err != nil { + return nil, err + } workerBind := optional(lookup, EnvDexWorkerBindAddress, defaultWorkerBindAddress) workerTarget := optional(lookup, EnvDexWorkerTarget, workerBind) config := &Config{ @@ -189,7 +195,8 @@ func load(lookup Lookup) (*Config, error) { APIKey: optional(lookup, EnvGroqAPIKey, ""), BaseURL: optional(lookup, EnvGroqBaseURL, ""), }, - RequestTimeout: defaultProviderRequestTimeout, + OpenAIDisableModelResponseStore: disableOpenAIResponseStore, + RequestTimeout: defaultProviderRequestTimeout, }, } if err := config.Validate(); err != nil { @@ -356,6 +363,18 @@ func optional(lookup Lookup, name EnvironmentVariable, fallback string) string { return strings.TrimSpace(value) } +func optionalBool(lookup Lookup, name EnvironmentVariable, fallback bool) (bool, error) { + value, found := lookup(string(name)) + if !found || strings.TrimSpace(value) == "" { + return fallback, nil + } + parsed, err := strconv.ParseBool(strings.TrimSpace(value)) + if err != nil { + return false, fmt.Errorf("%s must be a boolean", name) + } + return parsed, nil +} + func optionalPositiveInt64(lookup Lookup, name EnvironmentVariable, fallback int64) (int64, error) { value, found := lookup(string(name)) if !found || strings.TrimSpace(value) == "" { diff --git a/internal/config/config_test.go b/internal/config/config_test.go index 07d6643..eba769a 100644 --- a/internal/config/config_test.go +++ b/internal/config/config_test.go @@ -28,11 +28,40 @@ func TestLoadUsesValidatedDefaults(t *testing.T) { len(config.HTTP.AllowedOrigins) != 0 || config.Events.RecoveryLimit != defaultStreamRecoveryLimit || config.Dex.WorkerTarget != defaultWorkerBindAddress || - config.BlobCache.MaxBytes != defaultBlobCacheMaxBytes { + config.BlobCache.MaxBytes != defaultBlobCacheMaxBytes || + config.Providers.OpenAIDisableModelResponseStore { t.Fatalf("unexpected defaults: %#v", config) } } +func TestLoadDisablesOpenAIResponseStoreExplicitly(t *testing.T) { + t.Parallel() + values := map[string]string{string(EnvOpenAIDisableModelResponseStore): "true"} + config, err := load(func(name string) (string, bool) { + value, found := values[name] + return value, found + }) + if err != nil { + t.Fatal(err) + } + if !config.Providers.OpenAIDisableModelResponseStore { + t.Fatal("OpenAI response store disable flag = false, want true") + } +} + +func TestLoadRejectsInvalidOpenAIResponseStoreFlag(t *testing.T) { + t.Parallel() + _, err := load(func(name string) (string, bool) { + if name == string(EnvOpenAIDisableModelResponseStore) { + return "sometimes", true + } + return "", false + }) + if err == nil { + t.Fatal("load error = nil") + } +} + func TestLoadConfiguresBoundedStreamRecovery(t *testing.T) { t.Parallel() values := map[string]string{string(EnvStreamRecoveryLimit): "250"} diff --git a/internal/model/anthropic.go b/internal/model/anthropic.go index cfda7a1..5f7e7bf 100644 --- a/internal/model/anthropic.go +++ b/internal/model/anthropic.go @@ -67,7 +67,12 @@ func NewAnthropicClient(credentials *CredentialStore, httpClient *http.Client, b } // Complete sends one typed Anthropic Messages request. -func (client *AnthropicClient) Complete(ctx context.Context, request agent.ModelRequest) (agent.ModelReply, error) { +func (client *AnthropicClient) Complete(ctx context.Context, request agent.ModelRequest) (agent.ModelReply, agent.ModelUsage, error) { + reply, err := client.complete(ctx, request) + return reply, agent.ModelUsage{}, err +} + +func (client *AnthropicClient) complete(ctx context.Context, request agent.ModelRequest) (agent.ModelReply, error) { if err := validateModelRequest(request, agent.ProviderAnthropic); err != nil { return agent.ModelReply{}, err } @@ -111,7 +116,12 @@ func (client *AnthropicClient) Complete(ctx context.Context, request agent.Model } // Summarize compacts application messages with Anthropic. -func (client *AnthropicClient) Summarize(ctx context.Context, request agent.SummarizeRequest) (string, error) { +func (client *AnthropicClient) Summarize(ctx context.Context, request agent.SummarizeRequest) (string, agent.ModelUsage, error) { + summary, err := client.summarize(ctx, request) + return summary, agent.ModelUsage{}, err +} + +func (client *AnthropicClient) summarize(ctx context.Context, request agent.SummarizeRequest) (string, error) { model := request.Config.Model if request.Config.CompactionModel != nil { model = *request.Config.CompactionModel diff --git a/internal/model/client.go b/internal/model/client.go index f75045f..acf3cd5 100644 --- a/internal/model/client.go +++ b/internal/model/client.go @@ -58,31 +58,31 @@ func NewClient( } // Complete routes one provider-qualified completion. -func (client *Client) Complete(ctx context.Context, request agent.ModelRequest) (agent.ModelReply, error) { +func (client *Client) Complete(ctx context.Context, request agent.ModelRequest) (agent.ModelReply, agent.ModelUsage, error) { provider, err := request.Config.Model.Provider() if err != nil { - return agent.ModelReply{}, err + return agent.ModelReply{}, agent.ModelUsage{}, err } adapter, err := client.adapter(provider) if err != nil { - return agent.ModelReply{}, err + return agent.ModelReply{}, agent.ModelUsage{}, err } return adapter.Complete(ctx, request) } // Summarize routes compaction to its explicitly configured provider. -func (client *Client) Summarize(ctx context.Context, request agent.SummarizeRequest) (string, error) { +func (client *Client) Summarize(ctx context.Context, request agent.SummarizeRequest) (string, agent.ModelUsage, error) { model := request.Config.Model if request.Config.CompactionModel != nil { model = *request.Config.CompactionModel } provider, err := model.Provider() if err != nil { - return "", err + return "", agent.ModelUsage{}, err } adapter, err := client.adapter(provider) if err != nil { - return "", err + return "", agent.ModelUsage{}, err } return adapter.Summarize(ctx, request) } diff --git a/internal/model/gemini.go b/internal/model/gemini.go index 48d2e34..7b5f1fe 100644 --- a/internal/model/gemini.go +++ b/internal/model/gemini.go @@ -64,7 +64,12 @@ func NewGeminiClient(credentials *CredentialStore, httpClient *http.Client, base } // Complete sends one typed Gemini generateContent request. -func (client *GeminiClient) Complete(ctx context.Context, request agent.ModelRequest) (agent.ModelReply, error) { +func (client *GeminiClient) Complete(ctx context.Context, request agent.ModelRequest) (agent.ModelReply, agent.ModelUsage, error) { + reply, err := client.complete(ctx, request) + return reply, agent.ModelUsage{}, err +} + +func (client *GeminiClient) complete(ctx context.Context, request agent.ModelRequest) (agent.ModelReply, error) { if err := validateModelRequest(request, agent.ProviderGemini); err != nil { return agent.ModelReply{}, err } @@ -116,7 +121,12 @@ func (client *GeminiClient) Complete(ctx context.Context, request agent.ModelReq } // Summarize compacts application messages with Gemini. -func (client *GeminiClient) Summarize(ctx context.Context, request agent.SummarizeRequest) (string, error) { +func (client *GeminiClient) Summarize(ctx context.Context, request agent.SummarizeRequest) (string, agent.ModelUsage, error) { + summary, err := client.summarize(ctx, request) + return summary, agent.ModelUsage{}, err +} + +func (client *GeminiClient) summarize(ctx context.Context, request agent.SummarizeRequest) (string, error) { model := request.Config.Model if request.Config.CompactionModel != nil { model = *request.Config.CompactionModel diff --git a/internal/model/groq.go b/internal/model/groq.go index 1892208..2930bd8 100644 --- a/internal/model/groq.go +++ b/internal/model/groq.go @@ -67,7 +67,12 @@ func NewGroqClient(credentials *CredentialStore, httpClient *http.Client, baseUR } // Complete sends one typed Groq completion and writes its complete visible text once. -func (client *GroqClient) Complete(ctx context.Context, request agent.ModelRequest) (agent.ModelReply, error) { +func (client *GroqClient) Complete(ctx context.Context, request agent.ModelRequest) (agent.ModelReply, agent.ModelUsage, error) { + reply, err := client.complete(ctx, request) + return reply, agent.ModelUsage{}, err +} + +func (client *GroqClient) complete(ctx context.Context, request agent.ModelRequest) (agent.ModelReply, error) { if err := validateModelRequest(request, agent.ProviderGroq); err != nil { return agent.ModelReply{}, err } @@ -138,7 +143,12 @@ func (client *GroqClient) Complete(ctx context.Context, request agent.ModelReque } // Summarize compacts application messages with Groq. -func (client *GroqClient) Summarize(ctx context.Context, request agent.SummarizeRequest) (string, error) { +func (client *GroqClient) Summarize(ctx context.Context, request agent.SummarizeRequest) (string, agent.ModelUsage, error) { + summary, err := client.summarize(ctx, request) + return summary, agent.ModelUsage{}, err +} + +func (client *GroqClient) summarize(ctx context.Context, request agent.SummarizeRequest) (string, error) { model := request.Config.Model if request.Config.CompactionModel != nil { model = *request.Config.CompactionModel diff --git a/internal/model/mock.go b/internal/model/mock.go index 76334a2..9b8a73f 100644 --- a/internal/model/mock.go +++ b/internal/model/mock.go @@ -40,7 +40,24 @@ func NewMockClient() *MockClient { } // Complete implements the deterministic mock/dex command language. -func (*MockClient) Complete(ctx context.Context, request agent.ModelRequest) (agent.ModelReply, error) { +func (client *MockClient) Complete(ctx context.Context, request agent.ModelRequest) (agent.ModelReply, agent.ModelUsage, error) { + reply, err := client.complete(ctx, request) + if err != nil { + return agent.ModelReply{}, agent.ModelUsage{}, err + } + inputTokens := int64(estimatedTokens(request.Messages)) + outputTokens := int64(len(reply.Content) / 4) + if outputTokens == 0 && reply.Content != "" { + outputTokens = 1 + } + return reply, agent.ModelUsage{ + InputTokens: inputTokens, + OutputTokens: outputTokens, + TotalTokens: inputTokens + outputTokens, + }, nil +} + +func (*MockClient) complete(ctx context.Context, request agent.ModelRequest) (agent.ModelReply, error) { if request.WriteAssistant == nil || request.WriteReasoning == nil || request.WriteActivity == nil { return agent.ModelReply{}, errors.New("mock model writers are required") } @@ -282,7 +299,24 @@ func (*MockClient) Complete(ctx context.Context, request agent.ModelRequest) (ag } // Summarize creates a bounded local transcript. -func (*MockClient) Summarize(_ context.Context, request agent.SummarizeRequest) (string, error) { +func (client *MockClient) Summarize(ctx context.Context, request agent.SummarizeRequest) (string, agent.ModelUsage, error) { + summary, err := client.summarize(ctx, request) + if err != nil { + return "", agent.ModelUsage{}, err + } + inputTokens := int64(estimatedTokens(request.Messages)) + outputTokens := int64(len(summary) / 4) + if outputTokens == 0 && summary != "" { + outputTokens = 1 + } + return summary, agent.ModelUsage{ + InputTokens: inputTokens, + OutputTokens: outputTokens, + TotalTokens: inputTokens + outputTokens, + }, nil +} + +func (*MockClient) summarize(_ context.Context, request agent.SummarizeRequest) (string, error) { parts := make([]string, 0, len(request.Messages)+1) if request.PreviousSummary != "" { parts = append(parts, request.PreviousSummary) diff --git a/internal/model/mock_test.go b/internal/model/mock_test.go index dee023d..95292ec 100644 --- a/internal/model/mock_test.go +++ b/internal/model/mock_test.go @@ -61,7 +61,7 @@ func TestMockClientPlansAndWaitsDeterministically(t *testing.T) { } for _, test := range tests { t.Run(test.name, func(t *testing.T) { - reply, err := client.Complete(context.Background(), agent.ModelRequest{ + reply, _, err := client.Complete(context.Background(), agent.ModelRequest{ Config: agent.NewAgentConfig(), Messages: []agent.AgentMessage{{Role: agent.MessageRoleUser, Content: test.request}}, Tools: test.tools, @@ -94,11 +94,11 @@ func TestMockClientProducesStableCallIDsForRetries(t *testing.T) { WriteActivity: discardActivity, FlowID: agent.FlowID("stable-flow"), } - first, err := NewMockClient().Complete(context.Background(), request) + first, _, err := NewMockClient().Complete(context.Background(), request) if err != nil { t.Fatalf("first Complete() error = %v", err) } - second, err := NewMockClient().Complete(context.Background(), request) + second, _, err := NewMockClient().Complete(context.Background(), request) if err != nil { t.Fatalf("second Complete() error = %v", err) } @@ -110,7 +110,7 @@ func TestMockClientProducesStableCallIDsForRetries(t *testing.T) { func TestMockClientDoesNotRepeatFailureCommandAfterToolResult(t *testing.T) { toolName := agent.ToolNameSimulateFailure callID := agent.CallID("call-1") - reply, err := NewMockClient().Complete(context.Background(), agent.ModelRequest{ + reply, _, err := NewMockClient().Complete(context.Background(), agent.ModelRequest{ Config: agent.NewAgentConfig(), Messages: []agent.AgentMessage{ {Role: agent.MessageRoleUser, Content: "/tool-failure"}, @@ -134,7 +134,7 @@ func TestMockClientDoesNotRepeatFailureCommandAfterToolResult(t *testing.T) { func TestMockClientStreamsVisibleResponse(t *testing.T) { client := NewMockClient() var streamed strings.Builder - reply, err := client.Complete(context.Background(), agent.ModelRequest{ + reply, usage, err := client.Complete(context.Background(), agent.ModelRequest{ Config: agent.NewAgentConfig(), Messages: []agent.AgentMessage{{Role: agent.MessageRoleUser, Content: "hello"}}, WriteAssistant: func(chunk string) error { @@ -151,12 +151,15 @@ func TestMockClientStreamsVisibleResponse(t *testing.T) { if streamed.String() != reply.Content || reply.Content != "Local demo response: hello" { t.Fatalf("stream = %q, reply = %q", streamed.String(), reply.Content) } + if usage.InputTokens <= 0 || usage.OutputTokens <= 0 || usage.TotalTokens != usage.InputTokens+usage.OutputTokens { + t.Fatalf("usage = %#v", usage) + } } func TestMockClientStreamsProviderReasoningSummary(t *testing.T) { client := NewMockClient() var reasoning strings.Builder - reply, err := client.Complete(t.Context(), agent.ModelRequest{ + reply, _, err := client.Complete(t.Context(), agent.ModelRequest{ Config: agent.NewAgentConfig(), Messages: []agent.AgentMessage{{Role: agent.MessageRoleUser, Content: "/reason Checked constraints | Ready to proceed"}}, WriteAssistant: discardText, diff --git a/internal/model/openai.go b/internal/model/openai.go index dc9b6f7..5f7f95f 100644 --- a/internal/model/openai.go +++ b/internal/model/openai.go @@ -34,12 +34,13 @@ import ( const compactionInstruction = "Compact the conversation faithfully. Preserve decisions, user preferences, unresolved work, tool outcomes, identifiers, and facts needed by future turns." -// OpenAIClient implements the stateless OpenAI Responses API boundary. +// OpenAIClient implements the OpenAI Responses API boundary. type OpenAIClient struct { credentials *CredentialStore httpClient *http.Client baseURL string reasoningSummaryStreamingEnabled bool + disableModelResponseStore bool } var _ agent.ModelClient = (*OpenAIClient)(nil) @@ -48,6 +49,8 @@ var _ agent.ModelClient = (*OpenAIClient)(nil) type OpenAIClientConfig struct { // ReasoningSummaryStreamingEnabled defaults to false, is immutable after construction, and requests and forwards provider reasoning summaries when enabled. ReasoningSummaryStreamingEnabled bool + // DisableModelResponseStore explicitly sends store=false. The default explicitly sends store=true. + DisableModelResponseStore bool } // NewOpenAIClient constructs an adapter with an explicitly owned HTTP client. @@ -71,25 +74,26 @@ func NewOpenAIClient( httpClient: httpClient, baseURL: strings.TrimSpace(baseURL), reasoningSummaryStreamingEnabled: configuration.ReasoningSummaryStreamingEnabled, + disableModelResponseStore: configuration.DisableModelResponseStore, } } -// Complete streams one stateless Responses API turn. -func (client *OpenAIClient) Complete(ctx context.Context, request agent.ModelRequest) (agent.ModelReply, error) { +// Complete streams one Responses API turn. +func (client *OpenAIClient) Complete(ctx context.Context, request agent.ModelRequest) (agent.ModelReply, agent.ModelUsage, error) { if err := validateModelRequest(request, agent.ProviderOpenAI); err != nil { - return agent.ModelReply{}, err + return agent.ModelReply{}, agent.ModelUsage{}, err } modelName, err := request.Config.Model.ProviderModel() if err != nil { - return agent.ModelReply{}, err + return agent.ModelReply{}, agent.ModelUsage{}, err } input, err := openAIInput(request.Messages) if err != nil { - return agent.ModelReply{}, err + return agent.ModelReply{}, agent.ModelUsage{}, err } tools, err := openAITools(request.Tools) if err != nil { - return agent.ModelReply{}, err + return agent.ModelReply{}, agent.ModelUsage{}, err } params := responses.ResponseNewParams{ Instructions: openai.String(request.Config.SystemPrompt), @@ -100,7 +104,7 @@ func (client *OpenAIClient) Complete(ctx context.Context, request agent.ModelReq Include: []responses.ResponseIncludable{ responses.ResponseIncludableReasoningEncryptedContent, }, - Store: openai.Bool(false), + Store: openai.Bool(!client.disableModelResponseStore), ParallelToolCalls: openai.Bool(true), Tools: tools, } @@ -109,7 +113,7 @@ func (client *OpenAIClient) Complete(ctx context.Context, request agent.ModelReq } if request.ForcedTool != "" { if !hasDefinition(request.Tools, request.ForcedTool) { - return agent.ModelReply{}, fmt.Errorf("forced tool %q is not available", request.ForcedTool) + return agent.ModelReply{}, agent.ModelUsage{}, fmt.Errorf("forced tool %q is not available", request.ForcedTool) } params.ToolChoice.OfFunctionTool = &responses.ToolChoiceFunctionParam{ Name: string(request.ForcedTool), @@ -117,31 +121,31 @@ func (client *OpenAIClient) Complete(ctx context.Context, request agent.ModelReq } sdkClient := client.newSDKClient(request.FlowID) - stream := sdkClient.Responses.NewStreaming(ctx, params) + stream := sdkClient.Responses.NewStreaming(ctx, params, openAIRequestOptions(request.CallID)...) defer func() { _ = stream.Close() }() return consumeOpenAIStream(stream, request, client.reasoningSummaryStreamingEnabled) } -// Summarize compacts application history without persisting provider-side state. -func (client *OpenAIClient) Summarize(ctx context.Context, request agent.SummarizeRequest) (string, error) { +// Summarize compacts application history. +func (client *OpenAIClient) Summarize(ctx context.Context, request agent.SummarizeRequest) (string, agent.ModelUsage, error) { model := request.Config.Model if request.Config.CompactionModel != nil { model = *request.Config.CompactionModel } provider, err := model.Provider() if err != nil { - return "", err + return "", agent.ModelUsage{}, err } if provider != agent.ProviderOpenAI { - return "", fmt.Errorf("OpenAI adapter cannot summarize with provider %q", provider) + return "", agent.ModelUsage{}, fmt.Errorf("OpenAI adapter cannot summarize with provider %q", provider) } modelName, err := model.ProviderModel() if err != nil { - return "", err + return "", agent.ModelUsage{}, err } transcript, err := json.Marshal(request.Messages) if err != nil { - return "", fmt.Errorf("encode compaction transcript: %w", err) + return "", agent.ModelUsage{}, fmt.Errorf("encode compaction transcript: %w", err) } prompt := "Previous summary:\n" + request.PreviousSummary + "\n\nMessages:\n" + string(transcript) sdkClient := client.newSDKClient(request.FlowID) @@ -151,19 +155,26 @@ func (client *OpenAIClient) Summarize(ctx context.Context, request agent.Summari OfString: param.NewOpt(prompt), }, Model: modelName, - Store: openai.Bool(false), - }) + Store: openai.Bool(!client.disableModelResponseStore), + }, openAIRequestOptions(request.CallID)...) if err != nil { - return "", fmt.Errorf("OpenAI compaction request: %w", err) + return "", agent.ModelUsage{}, fmt.Errorf("OpenAI compaction request: %w", err) } if response.Status != responses.ResponseStatusCompleted { - return "", newOpenAIResponseError(response) + return "", agent.ModelUsage{}, newOpenAIResponseError(response) } summary := strings.TrimSpace(response.OutputText()) if summary == "" { - return "", errors.New("OpenAI compaction response was empty") + return "", agent.ModelUsage{}, errors.New("OpenAI compaction response was empty") + } + return summary, openAIUsage(response.Usage), nil +} + +func openAIRequestOptions(callID agent.CallID) []option.RequestOption { + if callID == "" { + return nil } - return summary, nil + return []option.RequestOption{option.WithHeader("Idempotency-Key", string(callID))} } // CountTokens returns a conservative local estimate used only to trigger compaction. @@ -284,26 +295,26 @@ func consumeOpenAIStream( stream openAIResponseStream, request agent.ModelRequest, reasoningSummaryStreamingEnabled bool, -) (agent.ModelReply, error) { +) (agent.ModelReply, agent.ModelUsage, error) { var content strings.Builder toolCalls := make(map[int64]agent.ToolCall) providerItems := make([]agent.ProviderContextItem, 0) - completed := false + var completed *responses.Response for stream.Next() { event := stream.Current() switch value := event.AsAny().(type) { case responses.ResponseTextDeltaEvent: if err := writeTextDelta(&content, value.Delta, request.WriteAssistant); err != nil { - return agent.ModelReply{}, err + return agent.ModelReply{}, agent.ModelUsage{}, err } case responses.ResponseRefusalDeltaEvent: if err := writeTextDelta(&content, value.Delta, request.WriteAssistant); err != nil { - return agent.ModelReply{}, err + return agent.ModelReply{}, agent.ModelUsage{}, err } case responses.ResponseReasoningSummaryTextDeltaEvent: if reasoningSummaryStreamingEnabled && value.Delta != "" { if err := request.WriteReasoning(value.Delta); err != nil { - return agent.ModelReply{}, fmt.Errorf("write reasoning summary: %w", err) + return agent.ModelReply{}, agent.ModelUsage{}, fmt.Errorf("write reasoning summary: %w", err) } } case responses.ResponseOutputItemDoneEvent: @@ -311,7 +322,7 @@ func consumeOpenAIStream( case responses.ResponseReasoningItem: encoded, err := encodeReasoningItem(item) if err != nil { - return agent.ModelReply{}, err + return agent.ModelReply{}, agent.ModelUsage{}, err } providerItems = append(providerItems, agent.ProviderContextItem{ Provider: agent.ProviderOpenAI, @@ -320,31 +331,31 @@ func consumeOpenAIStream( case responses.ResponseFunctionToolCall: call, err := openAIToolCall(item) if err != nil { - return agent.ModelReply{}, err + return agent.ModelReply{}, agent.ModelUsage{}, err } if _, duplicate := toolCalls[value.OutputIndex]; duplicate { - return agent.ModelReply{}, fmt.Errorf("OpenAI returned duplicate output index %d", value.OutputIndex) + return agent.ModelReply{}, agent.ModelUsage{}, fmt.Errorf("OpenAI returned duplicate output index %d", value.OutputIndex) } toolCalls[value.OutputIndex] = call if err := writeToolActivities([]agent.ToolCall{call}, request.WriteActivity); err != nil { - return agent.ModelReply{}, err + return agent.ModelReply{}, agent.ModelUsage{}, err } } case responses.ResponseCompletedEvent: - completed = true + completed = &value.Response case responses.ResponseFailedEvent: - return agent.ModelReply{}, newOpenAIResponseError(&value.Response) + return agent.ModelReply{}, agent.ModelUsage{}, newOpenAIResponseError(&value.Response) case responses.ResponseIncompleteEvent: - return agent.ModelReply{}, newOpenAIResponseError(&value.Response) + return agent.ModelReply{}, agent.ModelUsage{}, newOpenAIResponseError(&value.Response) case responses.ResponseErrorEvent: - return agent.ModelReply{}, fmt.Errorf("OpenAI stream error %q", value.Code) + return agent.ModelReply{}, agent.ModelUsage{}, fmt.Errorf("OpenAI stream error %q", value.Code) } } if err := stream.Err(); err != nil { - return agent.ModelReply{}, fmt.Errorf("read OpenAI response stream: %w", err) + return agent.ModelReply{}, agent.ModelUsage{}, fmt.Errorf("read OpenAI response stream: %w", err) } - if !completed { - return agent.ModelReply{}, errors.New("OpenAI response stream ended without a completed event") + if completed == nil { + return agent.ModelReply{}, agent.ModelUsage{}, errors.New("OpenAI response stream ended without a completed event") } ordered := make([]agent.ToolCall, 0, len(toolCalls)) for index := int64(0); len(ordered) < len(toolCalls); index++ { @@ -352,14 +363,24 @@ func consumeOpenAIStream( ordered = append(ordered, call) } if index > int64(len(toolCalls))+1024 { - return agent.ModelReply{}, errors.New("OpenAI returned invalid sparse tool indexes") + return agent.ModelReply{}, agent.ModelUsage{}, errors.New("OpenAI returned invalid sparse tool indexes") } } return agent.ModelReply{ Content: content.String(), ToolCalls: ordered, ProviderContextItems: providerItems, - }, nil + }, openAIUsage(completed.Usage), nil +} + +func openAIUsage(usage responses.ResponseUsage) agent.ModelUsage { + return agent.ModelUsage{ + InputTokens: usage.InputTokens, + CachedInputTokens: usage.InputTokensDetails.CachedTokens, + OutputTokens: usage.OutputTokens, + ReasoningTokens: usage.OutputTokensDetails.ReasoningTokens, + TotalTokens: usage.TotalTokens, + } } func writeTextDelta(content *strings.Builder, delta string, write agent.TextWriter) error { diff --git a/internal/model/openai_live_test.go b/internal/model/openai_live_test.go index 35df2b0..938c178 100644 --- a/internal/model/openai_live_test.go +++ b/internal/model/openai_live_test.go @@ -49,7 +49,7 @@ func TestLiveOpenAIResponses(t *testing.T) { config := agent.NewAgentConfig() config.Model = "openai/gpt-5-mini" config.SystemPrompt = "Answer concisely and accurately." - reply, err := client.Complete(ctx, agent.ModelRequest{ + reply, usage, err := client.Complete(ctx, agent.ModelRequest{ Config: config, Messages: []agent.AgentMessage{{ Role: agent.MessageRoleUser, @@ -66,6 +66,9 @@ func TestLiveOpenAIResponses(t *testing.T) { if err != nil { t.Fatal(err) } + if usage.TotalTokens <= 0 { + t.Fatalf("usage = %#v", usage) + } if strings.TrimSpace(reply.Content) == "" || reply.Content != streamed.String() { t.Fatal("OpenAI response was empty or differed from streamed content") } diff --git a/internal/model/openai_test.go b/internal/model/openai_test.go index 79f0e02..2fa03eb 100644 --- a/internal/model/openai_test.go +++ b/internal/model/openai_test.go @@ -30,7 +30,7 @@ import ( "github.com/superdurable/superagent/internal/agent" ) -func TestOpenAICompleteUsesStatelessResponsesAndSeparatesStreams(t *testing.T) { +func TestOpenAICompleteMapsUsageStoresByDefaultAndSeparatesStreams(t *testing.T) { t.Parallel() requestReceived := make(chan openAIRequestFixture, 1) httpClient := &http.Client{Transport: roundTripFunc(func(request *http.Request) (*http.Response, error) { @@ -40,6 +40,9 @@ func TestOpenAICompleteUsesStatelessResponsesAndSeparatesStreams(t *testing.T) { if request.Header.Get("Authorization") != "Bearer test-openai-key" { return fixtureResponse(request, http.StatusUnauthorized, "application/json", `{"error":{"message":"unexpected authorization"}}`), nil } + if request.Header.Get("Idempotency-Key") != "model-call-1" { + return fixtureResponse(request, http.StatusBadRequest, "application/json", `{"error":{"message":"unexpected idempotency key"}}`), nil + } var payload openAIRequestFixture if err := json.NewDecoder(request.Body).Decode(&payload); err != nil { return nil, fmt.Errorf("decode fixture request: %w", err) @@ -50,7 +53,7 @@ func TestOpenAICompleteUsesStatelessResponsesAndSeparatesStreams(t *testing.T) { sse(`{"type":"response.output_text.delta","delta":"I can help.","item_id":"msg-new","output_index":1,"content_index":0,"logprobs":[],"sequence_number":2}`), sse(`{"type":"response.output_item.done","output_index":0,"sequence_number":3,"item":{"id":"rs-new","type":"reasoning","summary":[{"type":"summary_text","text":"Checking constraints."}],"encrypted_content":"encrypted-new","status":"completed"}}`), sse(`{"type":"response.output_item.done","output_index":1,"sequence_number":4,"item":{"id":"fc-new","type":"function_call","call_id":"call-plan","name":"write_todos","arguments":"{\"todos\":[]}","status":"completed"}}`), - sse(`{"type":"response.completed","response":{"status":"completed"},"sequence_number":5}`), + sse(`{"type":"response.completed","response":{"status":"completed","usage":{"input_tokens":2,"input_tokens_details":{"cached_tokens":1,"cache_write_tokens":0},"output_tokens":4,"output_tokens_details":{"reasoning_tokens":3},"total_tokens":6}},"sequence_number":5}`), "data: [DONE]\n\n", }, "") return fixtureResponse(request, http.StatusOK, "text/event-stream", body), nil @@ -73,7 +76,7 @@ func TestOpenAICompleteUsesStatelessResponsesAndSeparatesStreams(t *testing.T) { assistantChunks := make([]string, 0) reasoningChunks := make([]string, 0) activity := make([]agent.AgentEvent, 0) - reply, err := client.Complete(context.Background(), agent.ModelRequest{ + reply, usage, err := client.Complete(context.Background(), agent.ModelRequest{ Config: agent.AgentConfig{ Model: "openai/gpt-5-mini", SystemPrompt: "Be helpful.", @@ -117,17 +120,21 @@ func TestOpenAICompleteUsesStatelessResponsesAndSeparatesStreams(t *testing.T) { return nil }, FlowID: flowID, + CallID: "model-call-1", }) if err != nil { t.Fatal(err) } + if usage.TotalTokens != 6 || usage.InputTokens != 2 || usage.CachedInputTokens != 1 || usage.OutputTokens != 4 || usage.ReasoningTokens != 3 { + t.Fatalf("usage = %#v", usage) + } payload := <-requestReceived if payload.Model != "gpt-5-mini" || payload.Instructions != "Be helpful." { t.Fatalf("unexpected request model/instructions: %#v", payload) } - if payload.Store || !payload.Stream || !payload.ParallelToolCalls { - t.Fatalf("request must be stateless, streaming, and allow parallel tool calls: %#v", payload) + if !payload.Store || !payload.Stream || !payload.ParallelToolCalls { + t.Fatalf("request must be stored, streaming, and allow parallel tool calls: %#v", payload) } if !slices.Equal(payload.Include, []string{"reasoning.encrypted_content"}) { t.Fatalf("include = %q", payload.Include) @@ -168,6 +175,57 @@ func TestOpenAICompleteUsesStatelessResponsesAndSeparatesStreams(t *testing.T) { } } +func TestOpenAIDisableModelResponseStoreAppliesToCompleteAndSummarize(t *testing.T) { + t.Parallel() + stores := make(chan bool, 2) + httpClient := &http.Client{Transport: roundTripFunc(func(request *http.Request) (*http.Response, error) { + var payload struct { + Store bool `json:"store"` + Stream bool `json:"stream"` + } + if err := json.NewDecoder(request.Body).Decode(&payload); err != nil { + return nil, err + } + stores <- payload.Store + if payload.Stream { + body := sse(`{"type":"response.output_text.delta","delta":"Done.","item_id":"msg","output_index":0,"content_index":0,"logprobs":[],"sequence_number":1}`) + + sse(`{"type":"response.completed","response":{"status":"completed","usage":{"input_tokens":1,"input_tokens_details":{"cached_tokens":0,"cache_write_tokens":0},"output_tokens":1,"output_tokens_details":{"reasoning_tokens":0},"total_tokens":2}},"sequence_number":2}`) + + "data: [DONE]\n\n" + return fixtureResponse(request, http.StatusOK, "text/event-stream", body), nil + } + body := `{"status":"completed","output":[{"type":"message","id":"msg","status":"completed","role":"assistant","content":[{"type":"output_text","text":"summary","annotations":[],"logprobs":[]}]}],"usage":{"input_tokens":3,"input_tokens_details":{"cached_tokens":1,"cache_write_tokens":0},"output_tokens":2,"output_tokens_details":{"reasoning_tokens":1},"total_tokens":5}}` + return fixtureResponse(request, http.StatusOK, "application/json", body), nil + })} + credentials := NewCredentialStore() + if err := credentials.SetDefaultAPIKey(agent.ProviderOpenAI, "test-openai-key"); err != nil { + t.Fatal(err) + } + client := NewOpenAIClient(credentials, httpClient, "https://api.openai.test", &OpenAIClientConfig{ + DisableModelResponseStore: true, + }) + config := agent.AgentConfig{Model: "openai/gpt-5-mini", SystemPrompt: "Be helpful."} + if _, _, err := client.Complete(t.Context(), agent.ModelRequest{ + Config: config, Messages: []agent.AgentMessage{{Role: agent.MessageRoleUser, Content: "Hello"}}, + WriteAssistant: discardText, WriteReasoning: discardText, WriteActivity: discardActivity, + FlowID: "flow-disable-store", CallID: "complete-call", + }); err != nil { + t.Fatal(err) + } + summary, usage, err := client.Summarize(t.Context(), agent.SummarizeRequest{ + Config: config, Messages: []agent.AgentMessage{{Role: agent.MessageRoleUser, Content: "Hello"}}, + FlowID: "flow-disable-store", CallID: "summary-call", + }) + if err != nil { + t.Fatal(err) + } + if summary != "summary" || usage.TotalTokens != 5 || usage.CachedInputTokens != 1 || usage.ReasoningTokens != 1 { + t.Fatalf("summary/usage = %q/%#v", summary, usage) + } + if first, second := <-stores, <-stores; first || second { + t.Fatalf("store values = %t, %t; want false, false", first, second) + } +} + func TestOpenAICompleteSkipsReasoningSummaryWhenStreamingIsDisabled(t *testing.T) { t.Parallel() requestReceived := make(chan openAIRequestFixture, 1) @@ -193,7 +251,7 @@ func TestOpenAICompleteSkipsReasoningSummaryWhenStreamingIsDisabled(t *testing.T } client := NewOpenAIClient(credentials, httpClient, "https://api.openai.test", &OpenAIClientConfig{}) reasoningChunks := make([]string, 0) - reply, err := client.Complete(context.Background(), agent.ModelRequest{ + reply, _, err := client.Complete(context.Background(), agent.ModelRequest{ Config: agent.AgentConfig{Model: "openai/gpt-5-mini", SystemPrompt: "Be helpful."}, Messages: []agent.AgentMessage{{Role: agent.MessageRoleUser, Content: "Say ready."}}, WriteAssistant: func(string) error { return nil }, @@ -253,7 +311,7 @@ func TestOpenAIDisablesSDKRetries(t *testing.T) { t.Fatal(err) } client := NewOpenAIClient(credentials, httpClient, "https://api.openai.test", &OpenAIClientConfig{}) - _, err := client.Complete(context.Background(), agent.ModelRequest{ + _, _, err := client.Complete(context.Background(), agent.ModelRequest{ Config: agent.AgentConfig{Model: "openai/gpt-5-mini", SystemPrompt: "Be helpful."}, Messages: []agent.AgentMessage{{Role: agent.MessageRoleUser, Content: "Hello"}}, WriteAssistant: func(string) error { return nil }, diff --git a/internal/model/providers_test.go b/internal/model/providers_test.go index 87c9916..4b3bdee 100644 --- a/internal/model/providers_test.go +++ b/internal/model/providers_test.go @@ -111,7 +111,7 @@ func assertProviderReply( t.Helper() assistant := make([]string, 0) activity := make([]agent.AgentEvent, 0) - reply, err := client.Complete(context.Background(), agent.ModelRequest{ + reply, _, err := client.Complete(context.Background(), agent.ModelRequest{ Config: agent.AgentConfig{Model: model, SystemPrompt: "Be helpful."}, Messages: []agent.AgentMessage{{Role: agent.MessageRoleUser, Content: "Check weather."}}, Tools: []agent.ToolDefinition{{ diff --git a/script/testdata/public-api-consumer/consumer_test.go b/script/testdata/public-api-consumer/consumer_test.go index 2be1af2..e765493 100644 --- a/script/testdata/public-api-consumer/consumer_test.go +++ b/script/testdata/public-api-consumer/consumer_test.go @@ -29,12 +29,12 @@ import ( type modelClient struct{} -func (modelClient) Complete(context.Context, agent.ModelRequest) (agent.ModelReply, error) { - return agent.ModelReply{}, nil +func (modelClient) Complete(context.Context, agent.ModelRequest) (agent.ModelReply, agent.ModelUsage, error) { + return agent.ModelReply{}, agent.ModelUsage{}, nil } -func (modelClient) Summarize(context.Context, agent.SummarizeRequest) (string, error) { - return "", nil +func (modelClient) Summarize(context.Context, agent.SummarizeRequest) (string, agent.ModelUsage, error) { + return "", agent.ModelUsage{}, nil } func (modelClient) CountTokens(agent.Model, []agent.AgentMessage) int {