feat(server): SessionStore.update() + per-turn cost tracking#117
Merged
Conversation
Add a partial-update method to the SessionStore ABC and implement on NullSessionStore, SqliteSessionStore, and PostgresSessionStore. First slice of #116 (and umbrella #104, Cost Tracking) -- the architecture doc states cost data piggybacks on SessionStore.update(), so this lays the storage plumbing before BaseAgent and Http wiring. Signature is keyword-only and additive (cost_data: dict | None = None today; future fields can be added without breaking callers). Returns True iff the session existed. cost_data is opaque JSON owned by BaseAgent; the store does a shallow merge so each per-turn write only carries the delta. SQLite stores it as TEXT and merges in Python; Postgres uses native JSONB || merge in a single statement. SqliteSessionStore.save() switches from INSERT OR REPLACE to ON CONFLICT DO UPDATE so cost_data survives saves of new messages -- the existing Postgres save() already had the right shape. Schema migration: SqliteSessionStore checks PRAGMA table_info and runs ALTER TABLE ADD COLUMN when the column is missing; Postgres uses ADD COLUMN IF NOT EXISTS. Existing databases pick up the column on first connect with no operator action. HttpSessionStore.update() raises NotImplementedError until the platform PATCH /v1/sessions/{id} endpoint lands (companion fipsagents-platform#2, agent-side wiring tracked separately as the next slice of #116). Refs: #116, #104 Assisted-by: Claude Code (Opus 4.7)
Replace the NotImplementedError stub on HttpSessionStore.update() with
the real platform call. Maps to PATCH /v1/sessions/{session_id} with
{"cost_data": {...}} as the partial-update body; 200 → True, 404 →
False without raising. cost_data=None short-circuits to exists() to
mirror the SQLite/Postgres update() behavior.
Companion to fipsagents-platform#2; the platform-side PATCH route
ships separately on that repo. Tests cover the wire shape against
httpx.MockTransport (method/path/body/auth/error mapping) and the
full agent → ASGI → SQLite round-trip via the in-process platform
fixture, asserting write-wins shallow-merge semantics.
The e2e round-trip reads cost_data straight out of the platform's
SQLite — GET /v1/sessions/{id} doesn't surface cost_data and
extending it is out of scope for this slice (belongs with the
platform-side PATCH endpoint's evolution).
Refs: #116, #104, fips-agents/fipsagents-platform#2
Assisted-by: Claude Code (Opus 4.7)
…e.update Wire the OpenAI chat-completions handler to extract token usage from each turn's StreamComplete event and accumulate it onto the session's cost_data via SessionStore.update(). Cost tracking is fully server-layer -- BaseAgent and astep_stream are unchanged. Adds a symmetric SessionStore.get_cost_data() reader so the accumulator can read the existing total before writing the cumulative one back. Implemented for Null/Sqlite/Postgres; HttpSessionStore raises NotImplementedError until the platform exposes a GET endpoint, in which case the accumulator falls back to a per-turn-delta write rather than crashing. update() failures are caught and logged so cost-tracking issues never break the chat response. Refs #116, #104. Assisted-by: Claude Code (Opus 4.7)
Records the SessionStore.update() / get_cost_data() additions and the server-layer per-turn cost accumulator that lands the first slice of #104. HTTP-backed deployments are flagged as last-write-wins until the platform exposes a GET endpoint (fipsagents-platform#4). Closes #116. Assisted-by: Claude Code (Opus 4.7)
This was referenced Apr 28, 2026
This file contains hidden or bidirectional Unicode text that may be interpreted or compiled differently than what appears below. To review, open the file in an editor that reveals hidden Unicode characters.
Learn more about bidirectional Unicode characters
Sign up for free
to join this conversation on GitHub.
Already have an account?
Sign in to comment
1 participant
Add this suggestion to a batch that can be applied as a single commit.This suggestion is invalid because no changes were made to the code.Suggestions cannot be applied while the pull request is closed.Suggestions cannot be applied while viewing a subset of changes.Only one suggestion per line can be applied in a batch.Add this suggestion to a batch that can be applied as a single commit.Applying suggestions on deleted lines is not supported.You must change the existing code in this line in order to create a valid suggestion.Outdated suggestions cannot be applied.This suggestion has been applied or marked resolved.Suggestions cannot be applied from pending reviews.Suggestions cannot be applied on multi-line comments.Suggestions cannot be applied while the pull request is queued to merge.Suggestion cannot be applied right now. Please check back later.
Summary
First slice of #104 (Cost Tracking) — adds
SessionStore.update()/get_cost_data()and wires the OpenAI server layer to record per-turn token usage on the session. Closes #116.The architecture doc states cost data piggybacks on
SessionStore.update(). This PR lands the storage plumbing + server wiring; pricing, budgets, and aggregation endpoints are deferred to follow-up issues.What landed
SessionStore.update(session_id, *, cost_data: dict | None = None) -> bool— partial update on the ABC. Shallow-merge per top-level key, write-wins. Implemented on Null / Sqlite / Postgres / Http (Http maps toPATCH /v1/sessions/{id}).SessionStore.get_cost_data(session_id) -> dict— symmetric reader so the server-side accumulator can read existing totals before writing cumulative ones back. Implemented on Null / Sqlite / Postgres; Http raisesNotImplementedErroruntil the platform exposes a GET endpoint (tracked as feat: GET /v1/sessions/{id}/cost_data for cumulative cost-tracking via HttpSessionStore fipsagents-platform#4).PRAGMA table_info+ALTER TABLE ADD COLUMN; Postgres viaADD COLUMN IF NOT EXISTS. No operator action required.SqliteSessionStore.save()switches fromINSERT OR REPLACEtoON CONFLICT DO UPDATE SET messages, updated_atsocost_datasurvives saves of new messages.OpenAIChatServer— both sync and streaming paths captureStreamMetrics.prompt_tokens/completion_tokensfrom the terminalStreamCompleteevent, accumulate cumulatively viaget_cost_data+update, and persist after each turn. Failures are caught and logged so cost-tracking never breaks the chat response.Companion PR
PATCH /v1/sessions/{session_id}route on the platform service. The agent-side e2e tests already pass against an in-process platform fixture; merging the platform PR is what makes the wire endpoint reachable in deployed clusters.Out of scope (filed as follow-ups)
BudgetEnforcerobserver with soft/hard limits +BudgetExceededErrorGET /v1/sessions/{id}/usageREST aggregation endpointtenant_id/session_idlabel dimensions onagent_tokens_totalPrometheus metricGET /v1/sessions/{id}/cost_dataso HTTP-backed deployments get cumulative semantics (today HTTP falls back to per-turn-delta writes)Test plan
pytest tests/test_sessions.py— 25/25 (8 new forupdate()+get_cost_data()+ migration)pytest tests/test_http_stores.py— 28/28 (4 new forupdate()PATCH wire shape, 1 new forget_cost_dataNotImplementedError)pytest tests/test_http_stores_e2e.py— 17/17 (3 new for full agent → ASGI → SQLite round-trip via in-process platform)pytest tests/test_server_openai.py— full suite (7 new for cumulative accumulation, no-session no-op, persist-failure isolation, NullSessionStore default)Refs #104, #116.
Assisted-by: Claude Code (Opus 4.7)