fix(gateway): preserve owner isolation when thread metadata is missing (#5484)

* fix(gateway): preserve owner isolation when thread metadata is missing

Follow-up to the #5448 review P1 (post-merge finding): owner_check=True
also authorizes threads whose meta row is missing (legacy compatibility)
or NULL-owner (shared/pre-auth data). _run_scope_user_id returned None
for every trusted internal caller, which dropped the only remaining
per-user filter on those threads and let an internal caller acting for
owner A read owner B's persisted runs.

_run_scope_user_id now takes the thread_id and consults the thread meta
store: when an existing meta row establishes ownership, the authorized
thread's runs are still read unfiltered (merged #5448 semantics,
including owner-header-less internal callers); when the meta row is
missing or NULL-owner, the filter falls back to the acting owner's raw
stamp (the exact value start_run writes) — or the synthetic "default"
identity without an owner header — so cross-user runs stay hidden.

Isolation coverage uses the real MemoryThreadMetaStore with no metadata
row (and a NULL-owner row) plus another user's persisted run: /runs and
/runs/page must be empty and /runs/{run_id} must 404 for internal
callers, while an established-ownership thread keeps the unfiltered
read.

* fix(gateway): gate run-scoped sub-resource reads for internal callers

Review follow-up on #5484: the P1 owner-isolation class remained
reachable through run-scoped sibling reads that apply no per-user filter
at all — /runs/{run_id}/messages, /events, /join, /stream and
/workspace-changes query by (thread_id, run_id) directly, so on
missing/NULL-owner threads an internal caller acting for owner A could
still read owner B's run content by id (verified 200 at the previous
head).

- Extract _thread_ownership_established (shared meta-row check) and add
  _require_run_visible_to_scope: for internal callers on threads without
  established ownership, the run's own user_id stamp must match the
  acting owner's raw value (or the legacy "default" stamp) or the read
  404s. Established-ownership threads and every non-internal caller keep
  their existing thread-scoped semantics.
- Wire the gate into join, stream, messages, events and
  workspace-changes; reword the now-stale messages comment to track the
  new scoping semantics.

Regression tests: sub-resource reads 404 for a mismatched internal
owner while the matching owner reads them normally, and the owner-less
fallback branch (synthetic "default" filter on missing-meta threads) is
pinned. Red confirmed against the pre-gate head.

* fix(gateway): gate cancel and artifact archive for internal callers

Review follow-up on #5484 round 2: POST /cancel resolved runs unscoped
(require_existing=True only closes the missing-meta case — NULL-owner
meta rows still pass), so an internal caller acting for a different
owner could interrupt another owner's active run on a shared thread
while /join and /stream were already gated. The archive manifest and
download pair likewise leaked the other owner's delivered-file count
and a 200-vs-409 delivery oracle on NULL-owner threads (missing-meta
threads were already denied by require_existing=True).

All three routes now call _require_run_visible_to_scope; its docstring
records the extended coverage. NULL-owner-thread regression tests pin:
a mismatched internal owner gets 404 from cancel, manifest and archive
download, while the acting owner reaches the real conflict path (409 on
a terminal run) and reads the manifest (file_count 2).

* fix(gateway): tolerate state-less request stand-ins in the scope helpers

The new owner-isolation gate and _run_scope_user_id read request.state
directly, which crashed the FakeRequest-based unit suites for the run
events, workspace-changes and scope endpoints (backend-unit-tests shards
1/2/4 on #5484). Read the state object defensively first: a request
without state is simply not an internal caller, so those paths keep
their pre-gate semantics.

* fix(gateway): scope the thread token-usage aggregate by owner

Review follow-up on #5484 round 4: GET /{thread_id}/token-usage called
aggregate_tokens_by_thread(thread_id) with no user filter at all, so on
missing/NULL-owner threads an internal caller acting for owner A read
owner B's spend, model names, run count and (with include_active=true)
live activity; the NULL-owner variant reached browser sessions too.
build_context_usage's latest-model lookup was unfiltered as well.

aggregate_tokens_by_thread gains an optional user_id (mirroring
list_by_thread: explicit None = unfiltered, AUTO resolves the contextvar)
in the memory store, the SQL repository and the store base;
build_context_usage/_resolve_thread_model_name thread the scope through
the latest-run lookup; the token-usage endpoint passes
_run_scope_user_id's value. Established-ownership threads aggregate
unfiltered as before; shared/missing-meta threads narrow to the acting
identity. Stale helper-test comment reworded after the #5482 merge
adaptation.

* test(gateway): pin the unfiltered aggregate on established-ownership threads

Review follow-up on #5484 round 5: the established-ownership branch of
the token-usage scoping (store receives user_id=None) was the only
unpinned half of the contract — the round-4 call-assertions never set
app.state.thread_store, so their None came from the user-less stand-in
path. test_token_usage_unfiltered_on_established_ownership_for_
internal_callers seeds an established meta row plus runs stamped by two
different identities and asserts the totals fold (166 = 111 + 55);
together with the isolation tests it now catches both failure modes
(always-stamp narrowing and always-None leak).
This commit is contained in:
xiaodu55 2026-09-18 09:39:14 +08:00 committed by GitHub
parent c24fd1e66f
commit 16f154f32b
No known key found for this signature in database
GPG Key ID: B5690EEEBB952194
7 changed files with 505 additions and 40 deletions

View File

@ -32,10 +32,14 @@ async def _load_checkpoint_messages(accessor: Any, config: dict[str, Any]) -> li
return list(values.get("messages") or []) return list(values.get("messages") or [])
async def _resolve_thread_model_name(run_store: Any, thread_id: str, app_config: Any) -> str | None: async def _resolve_thread_model_name(run_store: Any, thread_id: str, app_config: Any, user_id: str | None = None) -> str | None:
"""Prefer the latest run's model, then fall back to the first configured model.""" """Prefer the latest run's model, then fall back to the first configured model.
``user_id`` scopes the latest-run lookup (``None`` = unfiltered, matching
the thread-scoped aggregate semantics).
"""
try: try:
runs = await run_store.list_by_thread(thread_id, limit=1) runs = await run_store.list_by_thread(thread_id, limit=1, user_id=user_id)
except Exception: except Exception:
runs = [] runs = []
if runs: if runs:
@ -59,7 +63,7 @@ def build_context_usage_payload(*, token_count: int, max_context_tokens: int | N
} }
async def build_context_usage(request: Request, thread_id: str, run_store: Any) -> dict[str, Any] | None: async def build_context_usage(request: Request, thread_id: str, run_store: Any, user_id: str | None = None) -> dict[str, Any] | None:
"""Return approximate usage for the latest materialized thread checkpoint.""" """Return approximate usage for the latest materialized thread checkpoint."""
try: try:
app_config = get_config() app_config = get_config()
@ -79,7 +83,7 @@ async def build_context_usage(request: Request, thread_id: str, run_store: Any)
logger.warning("Failed to count context messages for thread %s", thread_id, exc_info=True) logger.warning("Failed to count context messages for thread %s", thread_id, exc_info=True)
return None return None
model_name = await _resolve_thread_model_name(run_store, thread_id, app_config) model_name = await _resolve_thread_model_name(run_store, thread_id, app_config, user_id=user_id)
model_config = app_config.get_model_config(model_name) if model_name else None model_config = app_config.get_model_config(model_name) if model_name else None
configured_window = getattr(model_config, "context_window", None) if model_config is not None else None configured_window = getattr(model_config, "context_window", None) if model_config is not None else None
max_context_tokens = int(configured_window) if configured_window else None max_context_tokens = int(configured_window) if configured_window else None

View File

@ -590,7 +590,7 @@ async def _find_target_run_id(
return source_run_id return source_run_id
run_mgr = get_run_manager(request) run_mgr = get_run_manager(request)
user_id = await _run_scope_user_id(request) user_id = await _run_scope_user_id(request, thread_id)
records = await run_mgr.list_by_thread(thread_id, user_id=user_id, limit=10) records = await run_mgr.list_by_thread(thread_id, user_id=user_id, limit=10)
fallback_record = next( fallback_record = next(
(record for record in records if record.status == RunStatus.success and _run_last_ai_matches_message(record, target_message)), (record for record in records if record.status == RunStatus.success and _run_last_ai_matches_message(record, target_message)),
@ -677,7 +677,7 @@ def _run_status_value(record: Any) -> str | None:
async def _require_successful_source_run(thread_id: str, run_id: str, request: Request) -> RunRecord: async def _require_successful_source_run(thread_id: str, run_id: str, request: Request) -> RunRecord:
run_mgr = get_run_manager(request) run_mgr = get_run_manager(request)
user_id = await _run_scope_user_id(request) user_id = await _run_scope_user_id(request, thread_id)
record = await run_mgr.get(run_id, user_id=user_id) record = await run_mgr.get(run_id, user_id=user_id)
if record is None: if record is None:
# The run-event journal is the authoritative lookup above. This fallback # The run-event journal is the authoritative lookup above. This fallback
@ -704,7 +704,7 @@ async def _find_interrupted_target_run_id(
return None return None
run_mgr = get_run_manager(request) run_mgr = get_run_manager(request)
user_id = await _run_scope_user_id(request) user_id = await _run_scope_user_id(request, thread_id)
record = await run_mgr.get(source_run_id, user_id=user_id) record = await run_mgr.get(source_run_id, user_id=user_id)
if record is None: if record is None:
records = await run_mgr.list_by_thread(thread_id, user_id=user_id, limit=20) records = await run_mgr.list_by_thread(thread_id, user_id=user_id, limit=20)
@ -1062,7 +1062,22 @@ def _parse_run_page_created_at(value: str) -> str:
return normalized return normalized
async def _run_scope_user_id(request: Request) -> str | None: async def _thread_ownership_established(request: Request, thread_id: str) -> bool:
"""Whether an existing meta row with a concrete owner covers ``thread_id``.
Missing rows (legacy compatibility) and NULL-owner rows (shared/pre-auth
data) do **not** establish ownership, even though ``owner_check=True``
still authorizes access to them.
"""
thread_store = getattr(request.app.state, "thread_store", None)
if thread_store is None:
return False
meta = await thread_store.get(thread_id, user_id=None)
meta_owner = meta.get("user_id") if isinstance(meta, dict) else getattr(meta, "user_id", None)
return meta is not None and bool(meta_owner)
async def _run_scope_user_id(request: Request, thread_id: str) -> str | None:
"""Resolve the data-filter id for run and message reads, not for authorization. """Resolve the data-filter id for run and message reads, not for authorization.
Thread visibility on these endpoints is already authorized by Thread visibility on these endpoints is already authorized by
@ -1071,10 +1086,17 @@ async def _run_scope_user_id(request: Request) -> str | None:
without an owner header, or the ``make_safe_user_id``-normalized owner without an owner header, or the ``make_safe_user_id``-normalized owner
otherwise — while ``start_run`` stamps run rows and run-event rows with otherwise — while ``start_run`` stamps run rows and run-event rows with
the raw trusted-owner value. Filtering by the authorization identity the raw trusted-owner value. Filtering by the authorization identity
therefore never matches the persisted rows (#5437), so internal callers therefore never matches the persisted rows (#5437).
read the authorized thread's runs, event-store messages, hidden-run
lookups, turn durations and feedback unfiltered; browser/API sessions Owner isolation (#5448 review P1): ``owner_check=True`` also authorizes
keep the per-user filter. threads whose meta row is missing (legacy compatibility) or NULL-owner
(shared/pre-auth data). On those, an unfiltered read would expose other
users' persisted runs to the acting owner's internal caller, so the
per-user filter is only dropped when the thread's meta row exists with an
established owner; otherwise the raw trusted owner — the exact value
``start_run`` stamps — is retained as the filter. Browser/API sessions
always keep the per-user filter. The thread token-usage aggregate is
scoped the same way.
Feedback note: an explicit ``None`` also skips the ``user_id`` WHERE in Feedback note: an explicit ``None`` also skips the ``user_id`` WHERE in
``FeedbackRepository``, so on shared/NULL-owner threads several users' ``FeedbackRepository``, so on shared/NULL-owner threads several users'
@ -1082,18 +1104,56 @@ async def _run_scope_user_id(request: Request) -> str | None:
/ ``list_by_run_ids`` order deterministically (latest wins, ``feedback_id`` / ``list_by_run_ids`` order deterministically (latest wins, ``feedback_id``
breaks ties) to keep that well-defined. breaks ties) to keep that well-defined.
""" """
user = getattr(request.state, "user", None) # Tolerate state-less request stand-ins used by focused unit tests.
if getattr(user, "system_role", None) == INTERNAL_SYSTEM_ROLE: state = getattr(request, "state", None)
user = getattr(state, "user", None)
if getattr(user, "system_role", None) != INTERNAL_SYSTEM_ROLE:
return await get_current_user(request)
if await _thread_ownership_established(request, thread_id):
return None return None
# Missing or NULL-owner meta row: ownership was never established, so the
# isolation boundary is the acting owner's raw stamp (the exact value
# start_run writes) — or, without an owner header, the synthetic
# "default" identity, which only matches legacy default-stamped rows.
owner = get_trusted_internal_owner_user_id(request)
if owner is not None:
return owner
return await get_current_user(request) return await get_current_user(request)
async def _require_run_visible_to_scope(run_id: str, thread_id: str, request: Request) -> None:
"""Gate run-scoped sub-resource reads and writes (events, messages, join,
stream, cancel, artifact archive).
These routes query or mutate by ``(thread_id, run_id)`` without a
per-user filter of their own. For trusted internal callers on threads
without established ownership, that let an internal caller acting for
owner A read or cancel owner B's run by id (#5448 review P1 follow-up).
The run's own stamp must therefore match the acting owner's raw value (or
the legacy ``"default"`` stamp); every other caller and every
established-ownership thread keeps its existing semantics.
"""
state = getattr(request, "state", None)
user = getattr(state, "user", None)
if getattr(user, "system_role", None) != INTERNAL_SYSTEM_ROLE:
return
if await _thread_ownership_established(request, thread_id):
return
scope = get_trusted_internal_owner_user_id(request) or "default"
record = await get_run_manager(request).get(run_id)
if record is None:
return
record_owner = getattr(record, "user_id", None) or "default"
if record.thread_id != thread_id or record_owner != scope:
raise HTTPException(status_code=404, detail=f"Run {run_id} not found")
@router.get("/{thread_id}/runs", response_model=list[RunResponse]) @router.get("/{thread_id}/runs", response_model=list[RunResponse])
@require_permission("runs", "read", owner_check=True) @require_permission("runs", "read", owner_check=True)
async def list_runs(thread_id: ThreadId, request: Request) -> list[RunResponse]: async def list_runs(thread_id: ThreadId, request: Request) -> list[RunResponse]:
"""List the newest runs for a thread (default 100, as a bare array).""" """List the newest runs for a thread (default 100, as a bare array)."""
run_mgr = get_run_manager(request) run_mgr = get_run_manager(request)
user_id = await _run_scope_user_id(request) user_id = await _run_scope_user_id(request, thread_id)
records = await run_mgr.list_by_thread(thread_id, user_id=user_id) records = await run_mgr.list_by_thread(thread_id, user_id=user_id)
return [_record_to_response(r) for r in records] return [_record_to_response(r) for r in records]
@ -1121,7 +1181,7 @@ async def list_runs_page(
before_created_at = _parse_run_page_created_at(before_created_at) before_created_at = _parse_run_page_created_at(before_created_at)
run_mgr = get_run_manager(request) run_mgr = get_run_manager(request)
user_id = await _run_scope_user_id(request) user_id = await _run_scope_user_id(request, thread_id)
records = await run_mgr.list_by_thread( records = await run_mgr.list_by_thread(
thread_id, thread_id,
user_id=user_id, user_id=user_id,
@ -1145,7 +1205,7 @@ async def list_runs_page(
async def get_run(thread_id: ThreadId, run_id: str, request: Request) -> RunResponse: async def get_run(thread_id: ThreadId, run_id: str, request: Request) -> RunResponse:
"""Get details of a specific run.""" """Get details of a specific run."""
run_mgr = get_run_manager(request) run_mgr = get_run_manager(request)
user_id = await _run_scope_user_id(request) user_id = await _run_scope_user_id(request, thread_id)
record = await run_mgr.get(run_id, user_id=user_id) record = await run_mgr.get(run_id, user_id=user_id)
if record is None or record.thread_id != thread_id: if record is None or record.thread_id != thread_id:
raise HTTPException(status_code=404, detail=f"Run {run_id} not found") raise HTTPException(status_code=404, detail=f"Run {run_id} not found")
@ -1172,6 +1232,7 @@ async def cancel_run(
durably notifies the owner when its lease is live, or takes over and durably notifies the owner when its lease is live, or takes over and
terminalizes the run when that lease has expired. terminalizes the run when that lease has expired.
""" """
await _require_run_visible_to_scope(run_id, thread_id, request)
run_mgr = get_run_manager(request) run_mgr = get_run_manager(request)
record = await run_mgr.get(run_id) record = await run_mgr.get(run_id)
if record is None or record.thread_id != thread_id: if record is None or record.thread_id != thread_id:
@ -1216,6 +1277,7 @@ async def cancel_run(
@require_permission("runs", "read", owner_check=True) @require_permission("runs", "read", owner_check=True)
async def join_run(thread_id: ThreadId, run_id: str, request: Request) -> StreamingResponse: async def join_run(thread_id: ThreadId, run_id: str, request: Request) -> StreamingResponse:
"""Join an existing run's SSE stream.""" """Join an existing run's SSE stream."""
await _require_run_visible_to_scope(run_id, thread_id, request)
run_mgr = get_run_manager(request) run_mgr = get_run_manager(request)
record = await run_mgr.get(run_id) record = await run_mgr.get(run_id)
if record is None or record.thread_id != thread_id: if record is None or record.thread_id != thread_id:
@ -1268,6 +1330,7 @@ async def _stream_existing_run(
""" """
require_cancel_permission_when_action(request, action) require_cancel_permission_when_action(request, action)
await _require_run_visible_to_scope(run_id, thread_id, request)
run_mgr = get_run_manager(request) run_mgr = get_run_manager(request)
record = await run_mgr.get(run_id) record = await run_mgr.get(run_id)
if record is None or record.thread_id != thread_id: if record is None or record.thread_id != thread_id:
@ -1368,10 +1431,11 @@ async def list_thread_messages(
after_seq: int | None = Query(default=None, ge=1), after_seq: int | None = Query(default=None, ge=1),
) -> list[dict]: ) -> list[dict]:
"""Return displayable messages for a thread (across all runs), with feedback attached.""" """Return displayable messages for a thread (across all runs), with feedback attached."""
# Resolve the data-filter id once (None for internal callers — same # Resolve the data-filter id once (None for internal callers on threads
# rationale as the runs endpoints above); it scopes the feedback query, # with established ownership — see `_run_scope_user_id`); it scopes the
# the hidden-run lookup, the event-store scan and turn-duration injection. # feedback query, the hidden-run lookup, the event-store scan and
user_id = await _run_scope_user_id(request) # turn-duration injection.
user_id = await _run_scope_user_id(request, thread_id)
run_mgr = get_run_manager(request) run_mgr = get_run_manager(request)
hidden_run_ids = await _default_history_hidden_run_ids(run_mgr, thread_id, user_id=user_id) hidden_run_ids = await _default_history_hidden_run_ids(run_mgr, thread_id, user_id=user_id)
messages, _ = await _scan_visible_thread_messages( messages, _ = await _scan_visible_thread_messages(
@ -1509,7 +1573,7 @@ async def list_thread_messages_page(
if "after_seq" in request.query_params: if "after_seq" in request.query_params:
raise HTTPException(status_code=422, detail="after_seq is not supported by this backward-only endpoint") raise HTTPException(status_code=422, detail="after_seq is not supported by this backward-only endpoint")
user_id = await _run_scope_user_id(request) user_id = await _run_scope_user_id(request, thread_id)
rows, has_more = await _scan_thread_message_page( rows, has_more = await _scan_thread_message_page(
thread_id, thread_id,
limit=limit, limit=limit,
@ -1539,6 +1603,7 @@ async def list_run_messages(
Response: { data: [...], has_more: bool } Response: { data: [...], has_more: bool }
""" """
await _require_run_visible_to_scope(run_id, thread_id, request)
event_store = get_run_event_store(request) event_store = get_run_event_store(request)
rows = await event_store.list_messages_by_run( rows = await event_store.list_messages_by_run(
thread_id, thread_id,
@ -1645,6 +1710,7 @@ async def get_run_artifact_archive_manifest(
request: Request, request: Request,
) -> ArtifactArchiveManifestResponse: ) -> ArtifactArchiveManifestResponse:
"""Return the verified terminal delivery count used by the archive.""" """Return the verified terminal delivery count used by the archive."""
await _require_run_visible_to_scope(run_id, thread_id, request)
presented_paths = await _archive_presented_paths(thread_id, run_id, request) presented_paths = await _archive_presented_paths(thread_id, run_id, request)
return ArtifactArchiveManifestResponse(file_count=len(dict.fromkeys(presented_paths))) return ArtifactArchiveManifestResponse(file_count=len(dict.fromkeys(presented_paths)))
@ -1657,6 +1723,7 @@ async def create_run_artifact_archive(
request: Request, request: Request,
) -> StreamingResponse: ) -> StreamingResponse:
"""Download the current contents of the files presented by one terminal run.""" """Download the current contents of the files presented by one terminal run."""
await _require_run_visible_to_scope(run_id, thread_id, request)
presented_paths = await _archive_presented_paths(thread_id, run_id, request) presented_paths = await _archive_presented_paths(thread_id, run_id, request)
raw_owner_user_id = get_trusted_internal_owner_user_id(request) raw_owner_user_id = get_trusted_internal_owner_user_id(request)
@ -1723,6 +1790,7 @@ async def list_run_events(
``task_id`` + ``after_seq`` let the subtask card page through one subagent ``task_id`` + ``after_seq`` let the subtask card page through one subagent
task's persisted steps without the run-wide ``limit`` truncating the tail (#3779). task's persisted steps without the run-wide ``limit`` truncating the tail (#3779).
""" """
await _require_run_visible_to_scope(run_id, thread_id, request)
event_store = get_run_event_store(request) event_store = get_run_event_store(request)
types = event_types.split(",") if event_types else None types = event_types.split(",") if event_types else None
events = await event_store.list_events( events = await event_store.list_events(
@ -1754,6 +1822,7 @@ async def get_run_workspace_changes(
include_diff: bool = Query(default=True), include_diff: bool = Query(default=True),
) -> dict: ) -> dict:
"""Return workspace/output file changes recorded for one run.""" """Return workspace/output file changes recorded for one run."""
await _require_run_visible_to_scope(run_id, thread_id, request)
event_store = get_run_event_store(request) event_store = get_run_event_store(request)
return await get_workspace_changes_response( return await get_workspace_changes_response(
event_store, event_store,
@ -1773,9 +1842,10 @@ async def thread_token_usage(
) -> ThreadTokenUsageResponse: ) -> ThreadTokenUsageResponse:
"""Thread-level token usage aggregation.""" """Thread-level token usage aggregation."""
run_store = get_run_store(request) run_store = get_run_store(request)
scope_user_id = await _run_scope_user_id(request, thread_id)
if include_active: if include_active:
agg = await run_store.aggregate_tokens_by_thread(thread_id, include_active=True) agg = await run_store.aggregate_tokens_by_thread(thread_id, include_active=True, user_id=scope_user_id)
else: else:
agg = await run_store.aggregate_tokens_by_thread(thread_id) agg = await run_store.aggregate_tokens_by_thread(thread_id, user_id=scope_user_id)
context_usage = await build_context_usage(request, thread_id, run_store) context_usage = await build_context_usage(request, thread_id, run_store, user_id=scope_user_id)
return ThreadTokenUsageResponse(thread_id=thread_id, context_usage=context_usage, **agg) return ThreadTokenUsageResponse(thread_id=thread_id, context_usage=context_usage, **agg)

View File

@ -495,7 +495,13 @@ class RunRepository(RunStore):
await session.execute(update(RunRow).where(RunRow.run_id == run_id, RunRow.status == "running").values(**values)) await session.execute(update(RunRow).where(RunRow.run_id == run_id, RunRow.status == "running").values(**values))
await session.commit() await session.commit()
async def aggregate_tokens_by_thread(self, thread_id: str, *, include_active: bool = False) -> dict[str, Any]: async def aggregate_tokens_by_thread(
self,
thread_id: str,
*,
include_active: bool = False,
user_id: str | None | _AutoSentinel = AUTO,
) -> dict[str, Any]:
"""Aggregate token usage for a thread. """Aggregate token usage for a thread.
``by_model`` is reduced in Python from each row's ``token_usage_by_model`` ``by_model`` is reduced in Python from each row's ``token_usage_by_model``
@ -513,6 +519,7 @@ class RunRepository(RunStore):
_completed = RunRow.status.in_(statuses) _completed = RunRow.status.in_(statuses)
_thread = RunRow.thread_id == thread_id _thread = RunRow.thread_id == thread_id
_run_operation = RunRow.operation_kind == "run" _run_operation = RunRow.operation_kind == "run"
resolved_user_id = resolve_user_id(user_id, method_name="RunRepository.aggregate_tokens_by_thread")
stmt = select( stmt = select(
RunRow.model_name, RunRow.model_name,
@ -524,6 +531,8 @@ class RunRepository(RunStore):
RunRow.middleware_tokens, RunRow.middleware_tokens,
RunRow.token_usage_by_model, RunRow.token_usage_by_model,
).where(_thread, _run_operation, _completed) ).where(_thread, _run_operation, _completed)
if resolved_user_id is not None:
stmt = stmt.where(RunRow.user_id == resolved_user_id)
async with self._sf() as session: async with self._sf() as session:
rows = (await session.execute(stmt)).all() rows = (await session.execute(stmt)).all()

View File

@ -299,7 +299,7 @@ class RunStore(abc.ABC):
pass pass
@abc.abstractmethod @abc.abstractmethod
async def aggregate_tokens_by_thread(self, thread_id: str, *, include_active: bool = False) -> dict[str, Any]: async def aggregate_tokens_by_thread(self, thread_id: str, *, include_active: bool = False, user_id: str | None = None) -> dict[str, Any]:
"""Aggregate token usage for completed runs in a thread. """Aggregate token usage for completed runs in a thread.
Returns a dict with keys: total_tokens, total_input_tokens, Returns a dict with keys: total_tokens, total_input_tokens,

View File

@ -252,12 +252,12 @@ class MemoryRunStore(RunStore):
results.sort(key=lambda r: r["created_at"]) results.sort(key=lambda r: r["created_at"])
return results return results
async def aggregate_tokens_by_thread(self, thread_id: str, *, include_active: bool = False) -> dict[str, Any]: async def aggregate_tokens_by_thread(self, thread_id: str, *, include_active: bool = False, user_id: str | None = None) -> dict[str, Any]:
statuses = ("success", "error", "running") if include_active else ("success", "error") statuses = ("success", "error", "running") if include_active else ("success", "error")
# Use the thread index for an O(runs-in-thread) lookup instead of # Use the thread index for an O(runs-in-thread) lookup instead of
# scanning every run in the process (mirrors ``list_by_thread``). # scanning every run in the process (mirrors ``list_by_thread``).
run_ids = self._runs_by_thread.get(thread_id) or () run_ids = self._runs_by_thread.get(thread_id) or ()
completed = [run for run_id in run_ids if (run := self._runs.get(run_id)) is not None and run.get("operation_kind", "run") == "run" and run.get("status") in statuses] completed = [run for run_id in run_ids if (run := self._runs.get(run_id)) is not None and run.get("operation_kind", "run") == "run" and run.get("status") in statuses and (user_id is None or run.get("user_id") == user_id)]
by_model: dict[str, dict] = {} by_model: dict[str, dict] = {}
for r in completed: for r in completed:
usage_by_model = r.get("token_usage_by_model") or {} usage_by_model = r.get("token_usage_by_model") or {}

View File

@ -19,6 +19,7 @@ from uuid import UUID
import pytest import pytest
from fastapi import FastAPI, HTTPException, Request, Response from fastapi import FastAPI, HTTPException, Request, Response
from fastapi.testclient import TestClient from fastapi.testclient import TestClient
from langgraph.store.memory import InMemoryStore
from starlette.middleware.base import BaseHTTPMiddleware from starlette.middleware.base import BaseHTTPMiddleware
from app.gateway.auth.models import User from app.gateway.auth.models import User
@ -26,6 +27,7 @@ from app.gateway.auth_disabled import AUTH_SOURCE_INTERNAL, AUTH_SOURCE_SESSION
from app.gateway.authz import AuthContext, Permissions from app.gateway.authz import AuthContext, Permissions
from app.gateway.internal_auth import INTERNAL_OWNER_USER_ID_HEADER_NAME, get_internal_user from app.gateway.internal_auth import INTERNAL_OWNER_USER_ID_HEADER_NAME, get_internal_user
from app.gateway.routers import thread_runs from app.gateway.routers import thread_runs
from deerflow.persistence.thread_meta.memory import MemoryThreadMetaStore
from deerflow.runtime.events.store.memory import MemoryRunEventStore from deerflow.runtime.events.store.memory import MemoryRunEventStore
from deerflow.runtime.runs.manager import RunManager from deerflow.runtime.runs.manager import RunManager
from deerflow.runtime.runs.store.memory import MemoryRunStore from deerflow.runtime.runs.store.memory import MemoryRunStore
@ -62,11 +64,18 @@ class _ScopeAuthMiddleware(BaseHTTPMiddleware):
class _PermissiveThreadStore: class _PermissiveThreadStore:
"""Stands in for the thread store behind ``owner_check=True``.""" """Stands in for the thread store behind ``owner_check=True``.
The existing scope tests exercise the established-ownership path, so
``get`` reports an existing, owner-established meta row.
"""
async def check_access(self, _thread_id: str, _user_id: str, *, require_existing: bool = False) -> bool: async def check_access(self, _thread_id: str, _user_id: str, *, require_existing: bool = False) -> bool:
return True return True
async def get(self, _thread_id: str, *, user_id: str | None | object = None) -> dict | None:
return {"thread_id": THREAD_ID, "user_id": "established-owner"}
class _RecordingRunStore(MemoryRunStore): class _RecordingRunStore(MemoryRunStore):
"""Records the per-user filter identity each read resolves to. """Records the per-user filter identity each read resolves to.
@ -141,10 +150,12 @@ def _make_app(
run_store: MemoryRunStore, run_store: MemoryRunStore,
event_store: MemoryRunEventStore | None = None, event_store: MemoryRunEventStore | None = None,
feedback_repo: _RecordingFeedbackRepo | None = None, feedback_repo: _RecordingFeedbackRepo | None = None,
thread_store=None,
) -> TestClient: ) -> TestClient:
app = FastAPI() app = FastAPI()
app.add_middleware(_ScopeAuthMiddleware, user=user, auth_source=auth_source) app.add_middleware(_ScopeAuthMiddleware, user=user, auth_source=auth_source)
app.state.thread_store = _PermissiveThreadStore() app.state.thread_store = thread_store if thread_store is not None else _PermissiveThreadStore()
app.state.run_store = run_store
app.state.run_manager = RunManager(store=run_store) app.state.run_manager = RunManager(store=run_store)
if event_store is not None: if event_store is not None:
app.state.run_event_store = event_store app.state.run_event_store = event_store
@ -319,16 +330,294 @@ def test_browser_session_messages_keep_per_user_filter(mixed_owner_store: Memory
# --------------------------------------------------------------------------- # ---------------------------------------------------------------------------
# owner isolation on threads without established metadata (#5448 review P1)
# ---------------------------------------------------------------------------
def test_missing_thread_meta_keeps_owner_isolation_for_internal_callers() -> None:
"""owner_check also authorizes missing-meta (legacy shared) threads.
There, unfiltered reads would expose other users' persisted runs to the
acting owner's internal caller, so the raw trusted owner stays the filter
— the exact value ``start_run`` stamps on run rows (#5448 review P1).
"""
thread_store = MemoryThreadMetaStore(InMemoryStore()) # no metadata row at all
run_store = _RecordingRunStore()
_seed_run(run_store, "run-other-user", user_id=str(BROWSER_USER_ID), status="success")
client = _make_app(
user=_internal_user(OWNER_RAW),
auth_source=AUTH_SOURCE_INTERNAL,
run_store=run_store,
thread_store=thread_store,
)
with client:
listed = client.get(
f"/api/threads/{THREAD_ID}/runs",
headers={INTERNAL_OWNER_USER_ID_HEADER_NAME: OWNER_RAW},
)
page = client.get(
f"/api/threads/{THREAD_ID}/runs/page",
headers={INTERNAL_OWNER_USER_ID_HEADER_NAME: OWNER_RAW},
)
single = client.get(
f"/api/threads/{THREAD_ID}/runs/run-other-user",
headers={INTERNAL_OWNER_USER_ID_HEADER_NAME: OWNER_RAW},
)
assert listed.status_code == 200
assert listed.json() == []
assert page.status_code == 200
assert page.json()["data"] == []
assert single.status_code == 404
# The acting owner's own raw-stamped runs remain visible: seed one and
# confirm it comes back through the same endpoints.
owned_store = _RecordingRunStore()
_seed_run(owned_store, "run-own-owner-stamp", user_id=OWNER_RAW, status="success")
client = _make_app(
user=_internal_user(OWNER_RAW),
auth_source=AUTH_SOURCE_INTERNAL,
run_store=owned_store,
thread_store=thread_store,
)
with client:
own = client.get(
f"/api/threads/{THREAD_ID}/runs/run-own-owner-stamp",
headers={INTERNAL_OWNER_USER_ID_HEADER_NAME: OWNER_RAW},
)
assert own.status_code == 200
assert own.json()["run_id"] == "run-own-owner-stamp"
def test_null_owner_thread_meta_keeps_owner_isolation_for_internal_callers() -> None:
"""NULL-owner meta rows (shared/pre-auth data) isolate by raw owner too."""
thread_store = MemoryThreadMetaStore(InMemoryStore())
asyncio.run(
thread_store.create(
THREAD_ID,
assistant_id=None,
user_id=None, # shared / pre-auth: meta row exists with NULL owner
)
)
run_store = _RecordingRunStore()
_seed_run(run_store, "run-other-user", user_id=str(BROWSER_USER_ID), status="success")
client = _make_app(
user=_internal_user(OWNER_RAW),
auth_source=AUTH_SOURCE_INTERNAL,
run_store=run_store,
thread_store=thread_store,
)
with client:
listed = client.get(
f"/api/threads/{THREAD_ID}/runs",
headers={INTERNAL_OWNER_USER_ID_HEADER_NAME: OWNER_RAW},
)
single = client.get(
f"/api/threads/{THREAD_ID}/runs/run-other-user",
headers={INTERNAL_OWNER_USER_ID_HEADER_NAME: OWNER_RAW},
)
assert listed.status_code == 200
assert listed.json() == []
assert single.status_code == 404
def test_established_ownership_still_reads_thread_runs_unfiltered(mixed_owner_store: MemoryRunStore) -> None:
"""Established meta ownership keeps the #5437 unfiltered-read behavior."""
thread_store = MemoryThreadMetaStore(InMemoryStore())
asyncio.run(thread_store.create(THREAD_ID, assistant_id=None, user_id=OWNER_RAW))
run_store = _RecordingRunStore()
_seed_run(run_store, RUN_OWNER, user_id=OWNER_RAW, status="success")
_seed_run(run_store, "run-legacy-default", user_id="default", status="success")
client = _make_app(
user=_internal_user(OWNER_RAW),
auth_source=AUTH_SOURCE_INTERNAL,
run_store=run_store,
thread_store=thread_store,
)
with client:
response = client.get(
f"/api/threads/{THREAD_ID}/runs",
headers={INTERNAL_OWNER_USER_ID_HEADER_NAME: OWNER_RAW},
)
assert response.status_code == 200
assert {row["run_id"] for row in response.json()} == {RUN_OWNER, "run-legacy-default"}
def test_ownerless_internal_caller_default_filter_on_missing_meta() -> None:
"""Without an owner header the synthetic "default" identity is the filter.
Pins the owner-less fallback branch of ``_run_scope_user_id``: a run
stamped with another owner's raw id stays hidden on missing-meta threads.
"""
thread_store = MemoryThreadMetaStore(InMemoryStore()) # no meta row
run_store = _RecordingRunStore()
_seed_run(run_store, "run-owner-777", user_id=OWNER_RAW, status="success")
client = _make_app(
user=_internal_user(None),
auth_source=AUTH_SOURCE_INTERNAL,
run_store=run_store,
thread_store=thread_store,
)
with client:
response = client.get(f"/api/threads/{THREAD_ID}/runs")
assert response.status_code == 200
assert response.json() == []
def test_subresource_reads_stay_owner_isolated_without_meta() -> None:
"""Run-scoped sub-resources must respect the acting owner's stamp.
These reads query by ``(thread_id, run_id)`` with no per-user filter of
their own; on missing-meta threads an internal caller acting for owner A
could otherwise read owner B's run content by id (#5448 review P1
follow-up).
"""
thread_store = MemoryThreadMetaStore(InMemoryStore()) # no meta row
run_store = _RecordingRunStore()
_seed_run(run_store, "run-owner-777", user_id=OWNER_RAW, status="success")
event_store = MemoryRunEventStore()
_seed_message(event_store, "run-owner-777", "msg-owner-run")
stranger_headers = {INTERNAL_OWNER_USER_ID_HEADER_NAME: "feishu:owner-999"}
owner_headers = {INTERNAL_OWNER_USER_ID_HEADER_NAME: OWNER_RAW}
base = f"/api/threads/{THREAD_ID}/runs/run-owner-777"
stranger = _make_app(
user=_internal_user("feishu:owner-999"),
auth_source=AUTH_SOURCE_INTERNAL,
run_store=run_store,
event_store=event_store,
thread_store=thread_store,
)
owner_client = _make_app(
user=_internal_user(OWNER_RAW),
auth_source=AUTH_SOURCE_INTERNAL,
run_store=run_store,
event_store=event_store,
thread_store=thread_store,
)
with stranger:
assert stranger.get(base + "/messages", headers=stranger_headers).status_code == 404
assert stranger.get(base + "/events", headers=stranger_headers).status_code == 404
assert stranger.get(base + "/workspace-changes", headers=stranger_headers).status_code == 404
assert stranger.get(base + "/join", headers=stranger_headers).status_code == 404
with owner_client:
messages = owner_client.get(base + "/messages", headers=owner_headers)
events = owner_client.get(base + "/events", headers=owner_headers)
assert messages.status_code == 200
assert [row["content"]["id"] for row in messages.json()["data"]] == ["msg-owner-run"]
assert events.status_code == 200
assert any(event.get("run_id") == "run-owner-777" for event in events.json())
def test_null_owner_thread_gates_cancel_and_archive_for_internal_callers() -> None:
"""NULL-owner meta rows gate POST /cancel and the archive pair too.
The round-2 findings: cancel resolved runs unscoped (an interrupt-vs-join
inconsistency) and the archive manifest leaked the other owner's
delivered-file count plus a 200-vs-409 delivery oracle on shared threads.
"""
thread_store = MemoryThreadMetaStore(InMemoryStore())
asyncio.run(thread_store.create(THREAD_ID, assistant_id=None, user_id=None))
run_store = _RecordingRunStore()
_seed_run(run_store, "run-owner-777", user_id=OWNER_RAW, status="success")
event_store = MemoryRunEventStore()
stranger_headers = {INTERNAL_OWNER_USER_ID_HEADER_NAME: "feishu:owner-999"}
stranger = _make_app(
user=_internal_user("feishu:owner-999"),
auth_source=AUTH_SOURCE_INTERNAL,
run_store=run_store,
event_store=event_store,
thread_store=thread_store,
)
with stranger:
cancel = stranger.post(
f"/api/threads/{THREAD_ID}/runs/run-owner-777/cancel?action=interrupt",
headers=stranger_headers,
)
manifest = stranger.get(
f"/api/threads/{THREAD_ID}/runs/run-owner-777/artifacts/archive",
headers=stranger_headers,
)
archive = stranger.post(
f"/api/threads/{THREAD_ID}/runs/run-owner-777/artifacts/archive",
headers=stranger_headers,
)
assert cancel.status_code == 404
assert manifest.status_code == 404
assert archive.status_code == 404
def test_null_owner_thread_matching_owner_cancels_and_reads_manifest() -> None:
"""The acting owner keeps cancel and archive access on shared threads."""
thread_store = MemoryThreadMetaStore(InMemoryStore())
asyncio.run(thread_store.create(THREAD_ID, assistant_id=None, user_id=None))
run_store = _RecordingRunStore()
_seed_run(run_store, "run-owner-777", user_id=OWNER_RAW, status="success")
event_store = MemoryRunEventStore()
asyncio.run(
event_store.put(
thread_id=THREAD_ID,
run_id="run-owner-777",
event_type="run.delivery",
category="outputs",
content={"presented": 2, "by_tool": {"present_files": ["/mnt/user-data/outputs/a.txt", "/mnt/user-data/outputs/b.txt"]}},
)
)
owner_client = _make_app(
user=_internal_user(OWNER_RAW),
auth_source=AUTH_SOURCE_INTERNAL,
run_store=run_store,
event_store=event_store,
thread_store=thread_store,
)
with owner_client:
manifest = owner_client.get(
f"/api/threads/{THREAD_ID}/runs/run-owner-777/artifacts/archive",
headers={INTERNAL_OWNER_USER_ID_HEADER_NAME: OWNER_RAW},
)
cancel = owner_client.post(
f"/api/threads/{THREAD_ID}/runs/run-owner-777/cancel?action=interrupt",
headers={INTERNAL_OWNER_USER_ID_HEADER_NAME: OWNER_RAW},
)
assert manifest.status_code == 200
assert manifest.json() == {"file_count": 2}
# A terminal run cannot be cancelled again: the acting owner reaches the
# real conflict path instead of a 404 anti-enumeration answer.
assert cancel.status_code == 409
# edit/regenerate helper fallback paths (#5482) # edit/regenerate helper fallback paths (#5482)
# --------------------------------------------------------------------------- # ---------------------------------------------------------------------------
def _helper_request(*, user, auth_source: str, run_store, event_store): def _helper_request(*, user, auth_source: str, run_store, event_store, owner_header: str | None = None):
"""Minimal Request stand-in: the helpers only touch state and app.state.""" """Minimal Request stand-in: the helpers touch state, app.state and headers."""
app_state = SimpleNamespace(run_manager=RunManager(store=run_store), run_event_store=event_store) app_state = SimpleNamespace(run_manager=RunManager(store=run_store), run_event_store=event_store)
headers = {INTERNAL_OWNER_USER_ID_HEADER_NAME: owner_header} if owner_header else {}
return SimpleNamespace( return SimpleNamespace(
state=SimpleNamespace(user=user, auth_source=auth_source), state=SimpleNamespace(user=user, auth_source=auth_source),
app=SimpleNamespace(state=app_state), app=SimpleNamespace(state=app_state),
headers=headers,
) )
@ -342,12 +631,13 @@ def test_helper_fallback_paths_resolve_internal_caller_runs() -> None:
auth_source=AUTH_SOURCE_INTERNAL, auth_source=AUTH_SOURCE_INTERNAL,
run_store=store, run_store=store,
event_store=MemoryRunEventStore(), event_store=MemoryRunEventStore(),
owner_header=OWNER_RAW,
) )
# The owner-stamped interrupted run resolves through the raw owner stamp. # The owner-stamped interrupted run resolves through the raw owner stamp.
interrupted = asyncio.run(thread_runs._find_interrupted_target_run_id(THREAD_ID, {"additional_kwargs": {"run_id": RUN_OWNER}}, request)) interrupted = asyncio.run(thread_runs._find_interrupted_target_run_id(THREAD_ID, {"additional_kwargs": {"run_id": RUN_OWNER}}, request))
assert interrupted == RUN_OWNER assert interrupted == RUN_OWNER
assert store.get_user_ids[-1] is None assert store.get_user_ids[-1] == OWNER_RAW
# An interrupted run is not an editable source run, but the lookup itself # An interrupted run is not an editable source run, but the lookup itself
# must have reached it (409 for status, not for a missing record). # must have reached it (409 for status, not for a missing record).
@ -355,10 +645,11 @@ def test_helper_fallback_paths_resolve_internal_caller_runs() -> None:
asyncio.run(thread_runs._require_successful_source_run(THREAD_ID, RUN_OWNER, request)) asyncio.run(thread_runs._require_successful_source_run(THREAD_ID, RUN_OWNER, request))
assert exc.value.status_code == 409 assert exc.value.status_code == 409
assert "successful" in exc.value.detail assert "successful" in exc.value.detail
assert store.get_user_ids[-1] is None assert store.get_user_ids[-1] == OWNER_RAW
# Fallback scan without any event-store or kwargs anchors still scans the # Fallback scan without any event-store or kwargs anchors still scans the
# authorized thread unfiltered (and 409s on the miss). # authorized thread through the acting owner's raw-stamp scope (and 409s
# on the miss).
with pytest.raises(HTTPException) as exc2: with pytest.raises(HTTPException) as exc2:
asyncio.run( asyncio.run(
thread_runs._find_target_run_id( thread_runs._find_target_run_id(
@ -370,7 +661,7 @@ def test_helper_fallback_paths_resolve_internal_caller_runs() -> None:
) )
) )
assert exc2.value.status_code == 409 assert exc2.value.status_code == 409
assert store.list_by_thread_user_ids and store.list_by_thread_user_ids[-1] is None assert store.list_by_thread_user_ids and store.list_by_thread_user_ids[-1] == OWNER_RAW
def test_helper_fallback_paths_keep_per_user_filter_for_browser_sessions() -> None: def test_helper_fallback_paths_keep_per_user_filter_for_browser_sessions() -> None:
@ -396,3 +687,94 @@ def test_helper_fallback_paths_keep_per_user_filter_for_browser_sessions() -> No
asyncio.run(thread_runs._require_successful_source_run(THREAD_ID, RUN_OWNER, request)) asyncio.run(thread_runs._require_successful_source_run(THREAD_ID, RUN_OWNER, request))
assert exc.value.status_code == 409 assert exc.value.status_code == 409
assert store.get_user_ids[-1] == str(BROWSER_USER_ID) assert store.get_user_ids[-1] == str(BROWSER_USER_ID)
def test_token_usage_isolated_without_meta_for_internal_callers() -> None:
"""Token-usage aggregate honors the acting owner's raw stamp (#5484 r4)."""
thread_store = MemoryThreadMetaStore(InMemoryStore()) # no meta row
run_store = _RecordingRunStore()
_seed_run(run_store, "run-owner-777", user_id=OWNER_RAW, status="success")
_seed_run(run_store, "run-other-user", user_id=str(BROWSER_USER_ID), status="success")
run_store._runs["run-owner-777"]["token_usage_by_model"] = {"gpt-x": {"total_tokens": 111}}
run_store._runs["run-owner-777"]["total_tokens"] = 111
run_store._runs["run-other-user"]["token_usage_by_model"] = {"gpt-x": {"total_tokens": 999}}
run_store._runs["run-other-user"]["total_tokens"] = 999
client = _make_app(
user=_internal_user(OWNER_RAW),
auth_source=AUTH_SOURCE_INTERNAL,
run_store=run_store,
thread_store=thread_store,
)
with client:
response = client.get(
f"/api/threads/{THREAD_ID}/token-usage",
headers={INTERNAL_OWNER_USER_ID_HEADER_NAME: OWNER_RAW},
)
assert response.status_code == 200
body = response.json()
assert body["total_tokens"] == 111
assert body["total_runs"] == 1
def test_token_usage_narrowed_for_browser_sessions_without_meta() -> None:
"""Browser sessions on shared threads see only their own spend too."""
thread_store = MemoryThreadMetaStore(InMemoryStore()) # no meta row
run_store = _RecordingRunStore()
_seed_run(run_store, "run-owner-777", user_id=OWNER_RAW, status="success")
_seed_run(run_store, "run-browser", user_id=str(BROWSER_USER_ID), status="success")
run_store._runs["run-owner-777"]["token_usage_by_model"] = {"gpt-x": {"total_tokens": 111}}
run_store._runs["run-owner-777"]["total_tokens"] = 111
run_store._runs["run-browser"]["token_usage_by_model"] = {"gpt-x": {"total_tokens": 222}}
run_store._runs["run-browser"]["total_tokens"] = 222
client = _make_app(
user=_browser_user(),
auth_source=AUTH_SOURCE_SESSION,
run_store=run_store,
thread_store=thread_store,
)
with client:
response = client.get(f"/api/threads/{THREAD_ID}/token-usage")
assert response.status_code == 200
assert response.json()["total_tokens"] == 222
def test_token_usage_unfiltered_on_established_ownership_for_internal_callers() -> None:
"""Established meta ownership keeps the unfiltered aggregate.
Pins the other half of the scoping contract: on an established thread the
internal caller's aggregate folds runs stamped by different identities
(the store must receive ``user_id=None``), mirroring
``test_established_ownership_still_reads_thread_runs_unfiltered``.
"""
thread_store = MemoryThreadMetaStore(InMemoryStore())
asyncio.run(thread_store.create(THREAD_ID, assistant_id=None, user_id=OWNER_RAW))
run_store = _RecordingRunStore()
_seed_run(run_store, "run-owner-777", user_id=OWNER_RAW, status="success")
_seed_run(run_store, "run-legacy-default", user_id="default", status="success")
run_store._runs["run-owner-777"]["token_usage_by_model"] = {"gpt-x": {"total_tokens": 111}}
run_store._runs["run-owner-777"]["total_tokens"] = 111
run_store._runs["run-legacy-default"]["token_usage_by_model"] = {"gpt-x": {"total_tokens": 55}}
run_store._runs["run-legacy-default"]["total_tokens"] = 55
client = _make_app(
user=_internal_user(OWNER_RAW),
auth_source=AUTH_SOURCE_INTERNAL,
run_store=run_store,
thread_store=thread_store,
)
with client:
response = client.get(
f"/api/threads/{THREAD_ID}/token-usage",
headers={INTERNAL_OWNER_USER_ID_HEADER_NAME: OWNER_RAW},
)
assert response.status_code == 200
body = response.json()
assert body["total_tokens"] == 166 # both stamps fold when ownership is established

View File

@ -58,7 +58,7 @@ def test_thread_token_usage_returns_stable_shape(monkeypatch: pytest.MonkeyPatch
**_aggregate_result(), **_aggregate_result(),
"context_usage": None, "context_usage": None,
} }
run_store.aggregate_tokens_by_thread.assert_awaited_once_with("thread-1") run_store.aggregate_tokens_by_thread.assert_awaited_once_with("thread-1", user_id=None)
build_context_usage.assert_awaited_once() build_context_usage.assert_awaited_once()
@ -72,7 +72,7 @@ def test_thread_token_usage_can_include_active_runs(monkeypatch: pytest.MonkeyPa
response = client.get("/api/threads/thread-1/token-usage?include_active=true") response = client.get("/api/threads/thread-1/token-usage?include_active=true")
assert response.status_code == 200 assert response.status_code == 200
run_store.aggregate_tokens_by_thread.assert_awaited_once_with("thread-1", include_active=True) run_store.aggregate_tokens_by_thread.assert_awaited_once_with("thread-1", include_active=True, user_id=None)
def test_thread_token_usage_serializes_context_percentage(monkeypatch: pytest.MonkeyPatch) -> None: def test_thread_token_usage_serializes_context_percentage(monkeypatch: pytest.MonkeyPatch) -> None: