From 9fb59785b285d16d2f56b24abe7b60b60283542b Mon Sep 17 00:00:00 2001 From: Pavel Date: Sat, 13 Jun 2026 15:39:16 +0400 Subject: [PATCH] Second fixes batch --- application/agents/tool_executor.py | 50 ++++++++--- .../0018_tool_attempts_attribution.py | 26 ++++-- application/api/user/analytics/routes.py | 61 ++++++++++---- application/llm/handlers/base.py | 13 +-- .../storage/db/repositories/stack_logs.py | 30 ++++++- .../db/repositories/tool_call_attempts.py | 52 +++++++++--- frontend/src/utils/dateTimeUtils.ts | 7 +- .../db/backfill_tool_attempts_attribution.py | 40 ++------- .../agents/test_tool_executor_three_phase.py | 82 ++++++++++++++++++ tests/api/user/test_analytics.py | 83 ++++++++++++++++++- .../db/repositories/test_stack_logs.py | 33 ++++++++ 11 files changed, 389 insertions(+), 88 deletions(-) diff --git a/application/agents/tool_executor.py b/application/agents/tool_executor.py index 0d059a0d..ccfdeb98 100644 --- a/application/agents/tool_executor.py +++ b/application/agents/tool_executor.py @@ -35,8 +35,13 @@ def _record_proposed( agent_id: Optional[str] = None, ) -> bool: """Insert a ``proposed`` row; swallow infra failures so tool calls - still run when the journal is unreachable. Returns True iff the row - is now journaled (newly created or already present). + still run when the journal is unreachable. Returns True iff THIS call + created the row. + + A duplicate ``call_id`` (LLMs reuse "call_0"-style ids) hits + ``ON CONFLICT DO NOTHING`` and returns False: the existing row may + belong to another in-flight request, so callers must not then flip it + via ``_mark_failed`` / ``_mark_executed``. """ try: with db_session() as conn: @@ -60,7 +65,7 @@ def _record_proposed( call_id, extra={"alert": "tool_call_id_collision", "call_id": call_id}, ) - return True + return inserted except Exception: logger.exception("tool_call_attempts proposed write failed for %s", call_id) return False @@ -77,11 +82,14 @@ def _mark_executed( action_name: Optional[str] = None, arguments: Any = None, tool_id: Optional[str] = None, + user_id: Optional[str] = None, + agent_id: Optional[str] = None, ) -> None: """Flip the row to ``executed``. If ``proposed_ok`` is False (the proposed write failed earlier), upsert a fresh row in ``executed`` so the reconciler can still see the attempt — without this, the side - effect would be invisible to the journal. + effect would be invisible to the journal. Both paths are scoped to + the owning ``user_id`` so a reused ``call_id`` can't cross tenants. """ try: with db_session() as conn: @@ -92,6 +100,7 @@ def _mark_executed( result, message_id=message_id, artifact_id=artifact_id, + user_id=user_id, ) if updated: return @@ -105,15 +114,23 @@ def _mark_executed( tool_id=tool_id if tool_id and looks_like_uuid(tool_id) else None, message_id=message_id, artifact_id=artifact_id, + user_id=user_id, + agent_id=( + str(agent_id) + if agent_id and looks_like_uuid(str(agent_id)) + else None + ), ) except Exception: logger.exception("tool_call_attempts executed write failed for %s", call_id) -def _mark_failed(call_id: str, error: str) -> None: +def _mark_failed(call_id: str, error: str, *, user_id: Optional[str] = None) -> None: try: with db_session() as conn: - ToolCallAttemptsRepository(conn).mark_failed(call_id, error) + ToolCallAttemptsRepository(conn).mark_failed( + call_id, error, user_id=user_id + ) except Exception: logger.exception("tool_call_attempts failed-write failed for %s", call_id) @@ -538,7 +555,9 @@ class ToolExecutor: user_id=self.user, agent_id=self.agent_id, ): - _mark_failed(call_id, tool_call_data["result"]) + _mark_failed( + call_id, tool_call_data["result"], user_id=self.user + ) yield {"type": "tool_call", "data": {**tool_call_data, "status": "error"}} self.tool_calls.append(tool_call_data) return "Failed to parse tool call.", call_id @@ -572,7 +591,11 @@ class ToolExecutor: user_id=self.user, agent_id=self.agent_id, ): - _mark_failed(call_id, f"Tool with ID {tool_id} not found.") + _mark_failed( + call_id, + f"Tool with ID {tool_id} not found.", + user_id=self.user, + ) yield {"type": "tool_call", "data": {**tool_call_data, "status": "error"}} self.tool_calls.append(tool_call_data) return f"Tool with ID {tool_id} not found.", call_id @@ -614,7 +637,8 @@ class ToolExecutor: ) tool_call_data["result"] = error_message tool_call_data["arguments"] = {} - _mark_failed(call_id, error_message) + if proposed_ok: + _mark_failed(call_id, error_message, user_id=self.user) yield { "type": "tool_call", "data": {**tool_call_data, "status": "error"}, @@ -677,7 +701,8 @@ class ToolExecutor: }, ) tool_call_data["result"] = error_message - _mark_failed(call_id, error_message) + if proposed_ok: + _mark_failed(call_id, error_message, user_id=self.user) yield {"type": "tool_call", "data": {**tool_call_data, "status": "error"}} self.tool_calls.append(tool_call_data) return error_message, call_id @@ -697,7 +722,8 @@ class ToolExecutor: logger.debug(f"Executing tool: {action_name} with args: {call_args}") result = tool.execute_action(action_name, **parameters) except Exception as exc: - _mark_failed(call_id, str(exc)) + if proposed_ok: + _mark_failed(call_id, str(exc), user_id=self.user) raise get_artifact_id = ( @@ -741,6 +767,8 @@ class ToolExecutor: action_name=action_name, arguments=call_args, tool_id=tool_data.get("id"), + user_id=self.user, + agent_id=self.agent_id, ) stream_tool_call_data = { diff --git a/application/alembic/versions/0018_tool_attempts_attribution.py b/application/alembic/versions/0018_tool_attempts_attribution.py index 0758d5ea..851d62c2 100644 --- a/application/alembic/versions/0018_tool_attempts_attribution.py +++ b/application/alembic/versions/0018_tool_attempts_attribution.py @@ -11,6 +11,11 @@ backfill lives in ``scripts/db/backfill_tool_attempts_attribution.py``; until it runs, the analytics reader falls back to the parent message's user for unstamped rows. +The index is built ``CONCURRENTLY`` in an autocommit block: the +``ADD COLUMN`` ALTERs commit first (a fast metadata-only change), so the +O(table-size) index build never runs while their ACCESS EXCLUSIVE lock +is held and never blocks live journaling. + Revision ID: 0018_tool_attempts_attribution Revises: 0017_oidc_scim """ @@ -29,14 +34,23 @@ depends_on: Union[str, Sequence[str], None] = None def upgrade() -> None: op.execute("ALTER TABLE tool_call_attempts ADD COLUMN user_id TEXT;") op.execute("ALTER TABLE tool_call_attempts ADD COLUMN agent_id UUID;") - op.execute( - "CREATE INDEX tool_call_attempts_user_ts_idx " - "ON tool_call_attempts (user_id, attempted_at DESC) " - "WHERE user_id IS NOT NULL;" - ) + # CONCURRENTLY can't run inside a transaction; the autocommit block + # commits the ALTERs above (releasing their lock) before building. + # DROP IF EXISTS first so a failed concurrent build (which leaves an + # INVALID index) is cleaned up on re-run. + with op.get_context().autocommit_block(): + op.execute("DROP INDEX IF EXISTS tool_call_attempts_user_ts_idx;") + op.execute( + "CREATE INDEX CONCURRENTLY tool_call_attempts_user_ts_idx " + "ON tool_call_attempts (user_id, attempted_at DESC) " + "WHERE user_id IS NOT NULL;" + ) def downgrade() -> None: - op.execute("DROP INDEX IF EXISTS tool_call_attempts_user_ts_idx;") + with op.get_context().autocommit_block(): + op.execute( + "DROP INDEX CONCURRENTLY IF EXISTS tool_call_attempts_user_ts_idx;" + ) op.execute("ALTER TABLE tool_call_attempts DROP COLUMN IF EXISTS agent_id;") op.execute("ALTER TABLE tool_call_attempts DROP COLUMN IF EXISTS user_id;") diff --git a/application/api/user/analytics/routes.py b/application/api/user/analytics/routes.py index 0c6839f2..0d244c2f 100644 --- a/application/api/user/analytics/routes.py +++ b/application/api/user/analytics/routes.py @@ -143,19 +143,19 @@ class GetMessageAnalytics(Resource): 200, ) - # Count messages per bucket, filtered by the conversation's - # owner (user_id) and optionally the agent. The ``user_id`` - # filter is always applied post-cutover to prevent - # cross-tenant leakage on admin dashboards. Agent matching - # covers both shapes: external traffic stamps ``api_key``, - # owner chats stamp ``agent_id``. + # Count messages per bucket. When filtering by agent the + # owner-scoped lookup above already gates access, so the + # user clause is dropped (matching tokens / tools / logs): + # a shared agent's conversations carry the caller's + # user_id, and the owner should see that traffic on their + # own agent's dashboard. Agent matching covers both + # shapes: external traffic stamps ``api_key``, owner / + # shared chats stamp ``agent_id``. clauses = [ - "c.user_id = :user_id", "m.timestamp >= :start", "m.timestamp <= :end", ] params: dict = { - "user_id": user, "start": start_date, "end": end_date, "fmt": pg_fmt, @@ -167,6 +167,9 @@ class GetMessageAnalytics(Resource): ) params["api_key"] = api_key params["agent_pg_id"] = agent_pg_id + else: + clauses.append("c.user_id = :user_id") + params["user_id"] = user where = " AND ".join(clauses) sql = ( "SELECT to_char(m.timestamp AT TIME ZONE 'UTC', :fmt) AS bucket, " @@ -385,24 +388,29 @@ class GetFeedbackAnalytics(Resource): # the timestamp from the JSONB and cast it to timestamptz for # the range filter + bucket grouping. clauses = [ - "c.user_id = :user_id", "m.feedback IS NOT NULL", "(m.feedback->>'timestamp')::timestamptz >= :start", "(m.feedback->>'timestamp')::timestamptz <= :end", ] params: dict = { - "user_id": user, "start": start_date, "end": end_date, "fmt": pg_fmt, } if api_key_id: + # Owner-gated agent match (see GetMessageAnalytics): + # drop the user clause so shared-agent feedback is + # visible to the owner, consistent with the other + # per-agent charts. clauses.append( "(c.api_key = :api_key" " OR c.agent_id = CAST(:agent_pg_id AS uuid))" ) params["api_key"] = api_key params["agent_pg_id"] = agent_pg_id + else: + clauses.append("c.user_id = :user_id") + params["user_id"] = user where = " AND ".join(clauses) sql = ( "SELECT to_char(" @@ -480,11 +488,15 @@ class GetToolAnalytics(Resource): jsonify({"success": True, "tools": []}), 200 ) - # Exclude non-terminal rows: a stuck ``proposed`` attempt - # (stream died, approval never granted) would otherwise - # render as a phantom success (successful = calls - failures). + # Terminal rows only. ``proposed`` (pending) and + # ``executed`` (ran, not yet finalized) are non-terminal: + # counting them inflates ``calls`` and — since the client + # computes successful = calls - failures — renders them as + # phantom successes that later flip to failures when the + # reconciler escalates a stuck row. ``confirmed`` is the + # only success state; ``failed`` the only failure. clauses = [ - "t.status <> 'proposed'", + "t.status IN ('confirmed', 'failed')", "t.attempted_at >= :start", "t.attempted_at <= :end", ] @@ -612,19 +624,24 @@ class GetScheduleAnalytics(Resource): # started/scheduled for runs that never got that far). ts = "COALESCE(r.finished_at, r.started_at, r.scheduled_for)" clauses = [ - "r.user_id = :user_id", f"{ts} >= :start", f"{ts} <= :end", ] params: dict = { - "user_id": user, "start": start_date, "end": end_date, "fmt": pg_fmt, } if api_key_id: + # Owner-gated agent match: drop the user clause so a + # shared agent's runs (created by callers under their + # own user_id via the scheduler tool) are visible to + # the owner, consistent with the per-agent timeline. clauses.append("r.agent_id = CAST(:agent_id AS uuid)") params["agent_id"] = agent_pg_id + else: + clauses.append("r.user_id = :user_id") + params["user_id"] = user where = " AND ".join(clauses) # The worker writes ``success`` / ``failed`` / ``timeout`` / # ``skipped`` (scheduler_worker.py); ``completed`` is kept in @@ -793,8 +810,11 @@ class GetUserLogs(Resource): "COALESCE(s.endpoint, '') NOT IN ('webhook', 'schedule')", "s.api_key = :api_key", ] + # Owner-gated agent match: drop the user clause so a + # shared agent's runs (stamped with the caller's + # user_id by the scheduler tool) appear on the owner's + # per-agent timeline, consistent with the chat branch. schedule_where = [ - "r.user_id = :user_id", "r.status IN ('success', 'completed', 'failed', 'timeout', 'skipped')", "r.agent_id = CAST(:agent_pg_id AS uuid)", ] @@ -991,10 +1011,15 @@ class GetUserLogs(Resource): ) ) + # ``ev.id`` is a unique per-branch tiebreaker: equal + # timestamps (transaction-stable ``now()`` makes ties + # routine) would otherwise sort non-deterministically + # across page queries, duplicating a row on one page and + # dropping its sibling from the next under OFFSET paging. sql = ( "SELECT * FROM (" + " UNION ALL ".join(branch_sqls) - + ") ev ORDER BY ev.timestamp DESC" + + ") ev ORDER BY ev.timestamp DESC, ev.id DESC" " LIMIT :limit OFFSET :offset" ) diff --git a/application/llm/handlers/base.py b/application/llm/handlers/base.py index c5963c54..6d325815 100644 --- a/application/llm/handlers/base.py +++ b/application/llm/handlers/base.py @@ -950,7 +950,7 @@ class LLMHandler(ABC): _record_proposed, ) - _record_proposed( + if _record_proposed( pause_info["call_id"], pause_info["tool_name"], pause_info["action_name"], @@ -959,11 +959,12 @@ class LLMHandler(ABC): message_id=agent.tool_executor.message_id, user_id=agent.tool_executor.user, agent_id=agent.tool_executor.agent_id, - ) - _mark_failed( - pause_info["call_id"], - f"headless: {deny_reason}", - ) + ): + _mark_failed( + pause_info["call_id"], + f"headless: {deny_reason}", + user_id=agent.tool_executor.user, + ) yield { "type": "tool_call", "data": { diff --git a/application/storage/db/repositories/stack_logs.py b/application/storage/db/repositories/stack_logs.py index a3cc3413..d24ed42f 100644 --- a/application/storage/db/repositories/stack_logs.py +++ b/application/storage/db/repositories/stack_logs.py @@ -17,6 +17,32 @@ from application.storage.db.serialization import PGNativeJSONEncoder from sqlalchemy import Connection, text +_REDACTED = "[REDACTED]" + + +def _is_secret_key(key: str) -> bool: + k = key.lower() + return k.endswith("api_key") or "secret" in k or "password" in k + + +def _redact_secrets(obj): + """Strip provider/agent keys from ``stacks`` before they persist. + + The ``llm`` stack component is built from every public attribute of + the LLM instance, which includes ``api_key`` (the deployment provider + secret) and ``user_api_key``. The unified-logs endpoint returns + ``stacks`` verbatim, so these must never reach the table. + """ + if isinstance(obj, dict): + return { + k: (_REDACTED if isinstance(k, str) and _is_secret_key(k) + else _redact_secrets(v)) + for k, v in obj.items() + } + if isinstance(obj, list): + return [_redact_secrets(v) for v in obj] + return obj + class StackLogsRepository: """Postgres-backed replacement for Mongo ``stack_logs`` collection.""" @@ -54,7 +80,9 @@ class StackLogsRepository: "user_id": user_id, "api_key": api_key, "query": query, - "stacks": json.dumps(stacks or [], cls=PGNativeJSONEncoder), + "stacks": json.dumps( + _redact_secrets(stacks or []), cls=PGNativeJSONEncoder + ), "timestamp": timestamp, }, ) diff --git a/application/storage/db/repositories/tool_call_attempts.py b/application/storage/db/repositories/tool_call_attempts.py index 6326f0cf..0954d733 100644 --- a/application/storage/db/repositories/tool_call_attempts.py +++ b/application/storage/db/repositories/tool_call_attempts.py @@ -75,13 +75,22 @@ class ToolCallAttemptsRepository: tool_id: Optional[str] = None, message_id: Optional[str] = None, artifact_id: Optional[str] = None, + user_id: Optional[str] = None, + agent_id: Optional[str] = None, ) -> None: """Insert OR upgrade a row to ``executed`` — or ``confirmed`` when there is no ``message_id``, as in ``mark_executed``. Used as a fallback when ``record_proposed`` failed (DB outage) and the tool ran anyway — preserves the journal so the - reconciler can still see the attempt. + reconciler can still see the attempt. ``user_id`` / ``agent_id`` + attribute the synthesized row so it stays visible to per-user / + per-agent analytics. + + The ``ON CONFLICT`` upgrade is guarded to a still-``proposed`` row + owned by the same ``user_id``: ``call_id`` is the table-wide PK + and LLMs reuse ids, so without it a colliding write could upgrade + another request's (possibly another tenant's) row. """ result_payload: dict = {"result": result} if artifact_id: @@ -92,16 +101,18 @@ class ToolCallAttemptsRepository: """ INSERT INTO tool_call_attempts (call_id, tool_id, tool_name, action_name, arguments, - result, message_id, status) + result, message_id, user_id, agent_id, status) VALUES (:call_id, CAST(:tool_id AS uuid), :tool_name, :action_name, CAST(:arguments AS jsonb), CAST(:result AS jsonb), CAST(:message_id AS uuid), - :status) + :user_id, CAST(:agent_id AS uuid), :status) ON CONFLICT (call_id) DO UPDATE SET status = :status, result = EXCLUDED.result, message_id = COALESCE(EXCLUDED.message_id, tool_call_attempts.message_id) + WHERE tool_call_attempts.status = 'proposed' + AND tool_call_attempts.user_id IS NOT DISTINCT FROM EXCLUDED.user_id """ ), { @@ -112,6 +123,8 @@ class ToolCallAttemptsRepository: "arguments": json.dumps(arguments if arguments is not None else {}, cls=PGNativeJSONEncoder), "result": json.dumps(result_payload, cls=PGNativeJSONEncoder), "message_id": message_id, + "user_id": user_id, + "agent_id": agent_id, "status": status, }, ) @@ -123,6 +136,7 @@ class ToolCallAttemptsRepository: *, message_id: Optional[str] = None, artifact_id: Optional[str] = None, + user_id: Optional[str] = None, ) -> bool: """Flip ``proposed`` → ``executed``, or straight to ``confirmed`` when there is no ``message_id`` (a ``save_conversation=False`` @@ -131,6 +145,11 @@ class ToolCallAttemptsRepository: ``artifact_id`` (when present) is stored alongside ``result`` in the JSONB as audit data — the reconciler reads it for diagnostic alerts when escalating stuck rows to ``failed``. + + Guarded to a still-``proposed`` row (and, when ``user_id`` is + given, to one owned by that user): ``call_id`` is the table-wide + PK and LLMs reuse ids, so an unguarded update could flip another + request's already-terminal — or another tenant's — row. """ result_payload: dict = {"result": result} if artifact_id: @@ -148,24 +167,33 @@ class ToolCallAttemptsRepository: if message_id is not None: sql += ", message_id = CAST(:message_id AS uuid)" params["message_id"] = message_id - sql += " WHERE call_id = :call_id" + sql += " WHERE call_id = :call_id AND status = 'proposed'" + if user_id is not None: + sql += " AND user_id IS NOT DISTINCT FROM :user_id" + params["user_id"] = user_id result_proxy = self._conn.execute(text(sql), params) return result_proxy.rowcount > 0 - def mark_failed(self, call_id: str, error: str) -> bool: + def mark_failed( + self, call_id: str, error: str, *, user_id: Optional[str] = None + ) -> bool: """Flip ``proposed`` → ``failed`` with the exception text. The status guard matters: ``call_id`` is the table-wide PK and LLMs have been observed reusing ids ("call_0"-style). Without it, a duplicate id hitting an error path would flip an already ``executed``/``confirmed`` row — possibly another request's — - to ``failed``, and nothing ever repairs that. + to ``failed``, and nothing ever repairs that. When ``user_id`` is + given the update is additionally scoped to that owner, so a + colliding id from another tenant can't fail this row. """ - result = self._conn.execute( - text( - "UPDATE tool_call_attempts SET status = 'failed', error = :error " - "WHERE call_id = :call_id AND status = 'proposed'" - ), - {"call_id": call_id, "error": error}, + sql = ( + "UPDATE tool_call_attempts SET status = 'failed', error = :error " + "WHERE call_id = :call_id AND status = 'proposed'" ) + params: dict[str, Any] = {"call_id": call_id, "error": error} + if user_id is not None: + sql += " AND user_id IS NOT DISTINCT FROM :user_id" + params["user_id"] = user_id + result = self._conn.execute(text(sql), params) return result.rowcount > 0 diff --git a/frontend/src/utils/dateTimeUtils.ts b/frontend/src/utils/dateTimeUtils.ts index b2fb65e8..ce7a2c39 100644 --- a/frontend/src/utils/dateTimeUtils.ts +++ b/frontend/src/utils/dateTimeUtils.ts @@ -28,7 +28,12 @@ export function formatDate(dateString: string): string { minute: '2-digit', }); } else if (/^\d{4}-\d{2}-\d{2}$/.test(dateString)) { - const date = new Date(dateString); + // `new Date('YYYY-MM-DD')` parses as UTC midnight, so rendering it + // in a timezone west of UTC shows the previous day. Construct from + // parts so the date is interpreted in local time (matching the + // space-separated hourly buckets, which already parse as local). + const [year, month, day] = dateString.split('-').map(Number); + const date = new Date(year, month - 1, day); return date.toLocaleDateString('en-US', { month: 'short', day: 'numeric', diff --git a/scripts/db/backfill_tool_attempts_attribution.py b/scripts/db/backfill_tool_attempts_attribution.py index d588f93a..f60f2dbe 100644 --- a/scripts/db/backfill_tool_attempts_attribution.py +++ b/scripts/db/backfill_tool_attempts_attribution.py @@ -8,16 +8,14 @@ Tiers ----- 1. Parent message (high confidence). Rows with a ``message_id`` copy the message's ``user_id`` and the conversation's ``agent_id``. -2. Schedule-run window (medium confidence). Headless rows never had a - message; a run's tool calls always fall inside that run's - started/finished window, so copy attribution from ``schedule_runs`` — - but only when exactly one run window contains the attempt, to avoid - cross-user misattribution on overlapping runs. -Rows matching neither tier (e.g. pre-0018 webhook runs, or attempts -whose parent message was deleted) are left NULL on purpose: the -analytics reader treats unattributable rows as invisible rather than -guessing an owner. +Message-less rows (headless: scheduled / webhook runs, plus pre-0018 +parse-failure rows) are left NULL on purpose: there is no FK linking an +attempt to its run, so any inference from a schedule-run *time window* +would also catch webhook attempts and misattribute them to an unrelated +tenant whose run happened to span the same instant. The analytics reader +treats unattributable rows as invisible rather than guessing an owner, +and new headless rows are stamped at propose time by the executor. Usage:: @@ -58,27 +56,6 @@ _TIER1 = text( """ ) -# Tier 2: headless rows via an unambiguous schedule-run time window. -_TIER2 = text( - """ - UPDATE tool_call_attempts t - SET user_id = r.user_id, - agent_id = r.agent_id - FROM schedule_runs r - WHERE t.user_id IS NULL - AND t.message_id IS NULL - AND r.started_at IS NOT NULL - AND r.finished_at IS NOT NULL - AND t.attempted_at BETWEEN r.started_at AND r.finished_at - AND ( - SELECT COUNT(*) FROM schedule_runs r2 - WHERE r2.started_at IS NOT NULL - AND r2.finished_at IS NOT NULL - AND t.attempted_at BETWEEN r2.started_at AND r2.finished_at - ) = 1 - """ -) - _COUNT_NULL = text( "SELECT count(*) FROM tool_call_attempts WHERE user_id IS NULL" ) @@ -108,14 +85,13 @@ def main() -> int: before = conn.execute(_COUNT_NULL).scalar_one() t1 = conn.execute(_TIER1).rowcount or 0 - t2 = conn.execute(_TIER2).rowcount or 0 after = conn.execute(_COUNT_NULL).scalar_one() print(f"NULL user_id rows before: {before}") print(f" tier 1 (parent message): {t1}") - print(f" tier 2 (schedule-run window): {t2}") print(f"NULL user_id rows remaining: {after}") + print(" (message-less headless rows left NULL by design)") if args.apply: trans.commit() diff --git a/tests/agents/test_tool_executor_three_phase.py b/tests/agents/test_tool_executor_three_phase.py index 88b01226..127b8a60 100644 --- a/tests/agents/test_tool_executor_three_phase.py +++ b/tests/agents/test_tool_executor_three_phase.py @@ -268,6 +268,88 @@ class TestRepository: assert row["status"] == "confirmed" assert row["error"] is None + def test_mark_executed_guarded_to_proposed(self, pg_conn): + """A reused ``call_id`` must not flip an already-terminal row back + to executed (the status guard mirrors ``mark_failed``).""" + from application.storage.db.repositories.tool_call_attempts import ( + ToolCallAttemptsRepository, + ) + + repo = ToolCallAttemptsRepository(pg_conn) + repo.record_proposed("c-me", "tool", "act", {}) + assert repo.mark_failed("c-me", "boom") is True + # Late executed emission for the same id: no proposed row remains. + assert repo.mark_executed("c-me", {"out": "ok"}) is False + row = _select_attempt(pg_conn, "c-me") + assert row["status"] == "failed" + assert row["result"] is None + + def test_mark_executed_scoped_to_owner(self, pg_conn): + """A colliding ``call_id`` from another tenant can't flip — or + read its result into — this user's proposed row.""" + from application.storage.db.repositories.tool_call_attempts import ( + ToolCallAttemptsRepository, + ) + + repo = ToolCallAttemptsRepository(pg_conn) + repo.record_proposed("call_0", "tool", "act", {}, user_id="victim") + # Attacker's request reuses the id and tries to land its result. + assert ( + repo.mark_executed("call_0", {"out": "leak"}, user_id="attacker") + is False + ) + row = _select_attempt(pg_conn, "call_0") + assert row["status"] == "proposed" + assert row["result"] is None + assert row["user_id"] == "victim" + + def test_mark_failed_scoped_to_owner(self, pg_conn): + from application.storage.db.repositories.tool_call_attempts import ( + ToolCallAttemptsRepository, + ) + + repo = ToolCallAttemptsRepository(pg_conn) + repo.record_proposed("call_0", "tool", "act", {}, user_id="victim") + assert repo.mark_failed("call_0", "boom", user_id="attacker") is False + row = _select_attempt(pg_conn, "call_0") + assert row["status"] == "proposed" + + def test_upsert_executed_stamps_attribution(self, pg_conn): + """The DB-outage fallback row must carry user/agent so it stays + visible to per-user / per-agent analytics (was born unattributed).""" + import uuid as _uuid + + from application.storage.db.repositories.tool_call_attempts import ( + ToolCallAttemptsRepository, + ) + + agent_id = str(_uuid.uuid4()) + repo = ToolCallAttemptsRepository(pg_conn) + repo.upsert_executed( + "c-up2", "tool", "act", {}, {"out": "ok"}, + user_id="u-up", agent_id=agent_id, + ) + row = _select_attempt(pg_conn, "c-up2") + assert row["user_id"] == "u-up" + assert str(row["agent_id"]) == agent_id + + def test_upsert_executed_wont_clobber_other_tenant(self, pg_conn): + """A colliding fallback upsert must not upgrade another tenant's + proposed row or overwrite its result.""" + from application.storage.db.repositories.tool_call_attempts import ( + ToolCallAttemptsRepository, + ) + + repo = ToolCallAttemptsRepository(pg_conn) + repo.record_proposed("call_0", "tool", "act", {}, user_id="victim") + repo.upsert_executed( + "call_0", "tool", "act", {}, {"out": "leak"}, user_id="attacker" + ) + row = _select_attempt(pg_conn, "call_0") + assert row["status"] == "proposed" + assert row["result"] is None + assert row["user_id"] == "victim" + @pytest.mark.unit class TestDefaultToolJournaling: diff --git a/tests/api/user/test_analytics.py b/tests/api/user/test_analytics.py index 0ee49c0b..4fb91723 100644 --- a/tests/api/user/test_analytics.py +++ b/tests/api/user/test_analytics.py @@ -31,12 +31,13 @@ def _patch_analytics_db(conn): def _seed_conversation_with_messages( pg_conn, user_id, *, count=3, api_key=None, feedback_text=None, + agent_id=None, ): from application.storage.db.repositories.conversations import ( ConversationsRepository, ) repo = ConversationsRepository(pg_conn) - conv = repo.create(user_id, name="t", api_key=api_key) + conv = repo.create(user_id, name="t", api_key=api_key, agent_id=agent_id) conv_id = str(conv["id"]) for i in range(count): repo.append_message(conv_id, {"prompt": f"p{i}", "response": f"r{i}"}) @@ -1104,3 +1105,83 @@ class TestUnifiedLogsWorkflowBranch: assert by_level["error"]["workflow_name"] == "My Flow" # Summary falls back to the workflow name when inputs lack a query. assert by_level["info"]["question"] == "My Flow" + + +class TestSharedAgentVisibility: + """Per-agent dashboards show the agent's full traffic, including + conversations a *shared* agent's callers create under their own + user_id. Messages / feedback / schedule must agree with the chat + timeline, which already exposes that traffic — they used to keep a + ``c.user_id``/``r.user_id`` clause that hid it (one screen, two + populations).""" + + def _owned_agent(self, pg_conn, owner): + from application.storage.db.repositories.agents import AgentsRepository + + return AgentsRepository(pg_conn).create( + owner, "shared-agent", "published", key="shared-key", + ) + + def test_message_analytics_includes_shared_caller(self, app, pg_conn): + from application.api.user.analytics.routes import GetMessageAnalytics + + agent = self._owned_agent(pg_conn, "owner-a") + # Caller B chats with A's shared agent: conversation.user_id = B, + # agent_id = A's agent. + _seed_conversation_with_messages( + pg_conn, "caller-b", count=3, agent_id=str(agent["id"]) + ) + response = _post_resource( + app, pg_conn, GetMessageAnalytics, "/api/get_message_analytics", + "owner-a", {"api_key_id": str(agent["id"])}, + ) + assert response.status_code == 200 + assert sum(response.json["messages"].values()) == 3 + + def test_feedback_analytics_includes_shared_caller(self, app, pg_conn): + from application.api.user.analytics.routes import GetFeedbackAnalytics + + agent = self._owned_agent(pg_conn, "owner-a") + _seed_conversation_with_messages( + pg_conn, "caller-b", count=2, agent_id=str(agent["id"]), + feedback_text="like", + ) + response = _post_resource( + app, pg_conn, GetFeedbackAnalytics, "/api/get_feedback_analytics", + "owner-a", {"api_key_id": str(agent["id"])}, + ) + assert response.status_code == 200 + totals = {"positive": 0, "negative": 0} + for bucket in response.json["feedback"].values(): + totals["positive"] += bucket["positive"] + totals["negative"] += bucket["negative"] + assert totals == {"positive": 1, "negative": 1} + + def test_schedule_analytics_includes_shared_caller(self, app, pg_conn): + helper = TestGetScheduleAnalytics() + agent = self._owned_agent(pg_conn, "owner-a") + # Run created by caller B against A's agent (scheduler tool stamps + # the caller's user_id). + helper._seed_run( + pg_conn, "caller-b", "success", agent_id=str(agent["id"]) + ) + response = helper._post( + app, pg_conn, "owner-a", {"api_key_id": str(agent["id"])} + ) + assert response.status_code == 200 + completed = sum(b["completed"] for b in response.json["runs"].values()) + assert completed == 1 + + def test_schedule_timeline_includes_shared_caller(self, app, pg_conn): + helper = TestGetScheduleAnalytics() + agent = self._owned_agent(pg_conn, "owner-a") + helper._seed_run( + pg_conn, "caller-b", "failed", agent_id=str(agent["id"]) + ) + response = _post_logs( + app, pg_conn, "owner-a", + {"api_key_id": str(agent["id"]), "event_type": "schedule"}, + ) + assert response.status_code == 200 + assert len(response.json["logs"]) == 1 + assert response.json["logs"][0]["level"] == "error" diff --git a/tests/storage/db/repositories/test_stack_logs.py b/tests/storage/db/repositories/test_stack_logs.py index 0a9218dc..99527ddc 100644 --- a/tests/storage/db/repositories/test_stack_logs.py +++ b/tests/storage/db/repositories/test_stack_logs.py @@ -50,3 +50,36 @@ class TestInsert: text("SELECT query FROM stack_logs WHERE activity_id = 'act-3'") ).fetchone() assert len(dict(row._mapping)["query"]) == 20000 + + def test_secrets_redacted_from_stacks(self, pg_conn): + # The ``llm`` stack component is built from every public attr of + # the LLM and includes the provider secret + the caller's key. + # These must never persist (the unified-logs endpoint returns + # stacks verbatim). + repo = _repo(pg_conn) + repo.insert( + activity_id="act-secret", + level="error", + stacks=[ + { + "component": "llm", + "data": { + "api_key": "sk-deployment-secret", + "user_api_key": "agent-key", + "OPENAI_API_KEY": "sk-env", + "model": "gpt-x", + "prompt_tokens": 42, + }, + } + ], + ) + row = pg_conn.execute( + text("SELECT stacks FROM stack_logs WHERE activity_id = 'act-secret'") + ).fetchone() + data = dict(row._mapping)["stacks"][0]["data"] + assert data["api_key"] == "[REDACTED]" + assert data["user_api_key"] == "[REDACTED]" + assert data["OPENAI_API_KEY"] == "[REDACTED]" + # Non-secret fields (incl. token *counts*) are untouched. + assert data["model"] == "gpt-x" + assert data["prompt_tokens"] == 42