feat(audit): record ai tool executions on every execution path#523
feat(audit): record ai tool executions on every execution path#523hamza-56 wants to merge 3 commits into
Conversation
hamza-56
force-pushed
the
feat/audit-tool-executions
branch
from
6fa85f0 to
351d59f
Compare
Tool executions were completely opaque: no record of who invoked which tool, when, with which inputs, from which model, or with what outcome (issue #501, part of epic #435). Every execution path now records a fail-closed tool.invoked event before the handler runs plus a tool.completed / tool.failed outcome event, paired by a shared invocation id, per ADR-0002: - Tool (MCP_TOOLS hook) audit-wraps its handler at construction, covering the FastMCP server and the chat tool registry at once - the chat provider loop stamps source=chat and model identity (provider, configured model, provider-reported version) - the RAG agent tools wrap explicitly and the agent runner stamps source=rag with best-effort model identity - a FastMCP middleware backstop records protocol-level failures that never reach a handler; the audit context middleware stamps source=mcp for the /ai mount - a completeness test pins the set of modules allowed to construct executable tools so new paths cannot appear unaudited Tool args are redacted before persistence (auth payloads and secret-named keys are replaced wholesale). No schema change: the audit_events AI-provenance columns shipped with #492.
hamza-56
force-pushed
the
feat/audit-tool-executions
branch
from
351d59f to
ac55801
Compare
hamza-56
requested review from
abdulrafey1,
regisb and
zaira-bibi
and removed request for
abdulrafey1,
regisb and
zaira-bibi
|
@claude-review |
|
Claude finished @hamza-56's task in 4m 17s —— View job Review: audit AI tool executions on every execution path
Reviewed all 28 files against A few points worth addressing before merge — none block the core design. Findings🟠 Medium — 🟠 Medium — anti-double-record guarantee lacks a direct test ( 🟡 Low — two committed transactions per tool call ( 🟡 Low / nit — stray contextvar set outside the middleware ( Things done well
Previous Review CommentsNo prior review threads on this PR. Overall: solid, mergeable work. Recommend addressing the |
![claude[bot]](https://avatars.githubusercontent.com/in/1236702?s=60&v=4){kind=small}
| tool=call, | ||
| model=model, | ||
| target=target, | ||
| error_detail=f"{type(exc).__name__}: {exc}", |
There was a problem hiding this comment.
Medium (security / defense-in-depth): error_detail is stored unredacted while tool_args is redacted.
tool.args goes through redact() in AuditEvent.from_event (models.py:116), but error_detail is persisted verbatim (models.py:154). The detail here is f"{type(exc).__name__}: {exc}", and several common exception types echo their input back in the message — Pydantic v2 ValidationError in particular renders input_value=... for each failing field. So a handler that raises while validating a credentials payload (or an httpx error carrying a tokenized URL) can leak the very secret that the auth/token redaction was designed to keep out of the store — and out of the canonical bytes.
Same path applies to the middleware backstop (sparkth/mcp/audit.py:50).
Consider scrubbing error_detail through the same secret filter (or capturing only type(exc).__name__ plus a length-bounded, secret-scrubbed message). At minimum, document that error_detail is not redacted so callers don't assume the whole row is safe.
There was a problem hiding this comment.
Fixed in 634b06f. Added scrub_error_detail in redaction.py, used by both the wrapper and the backstop: ValidationErrors are re-rendered without their inputs, secret-keyed assignments (token=..., 'api_key': ...) are masked with the same SECRET_KEYS set, and the result is capped at 1000 chars. Tests in test_redaction.py plus test_handler_error_detail_is_scrubbed end to end.
This comment was written with the assistance of an LLM (Claude).
| async def test_handler_level_failure_is_not_double_recorded(audit_events: AuditEventsFetcher) -> None: | ||
| """When the wrapped handler already recorded the failure, the middleware | ||
| backstop must stay silent: one invocation, one failure record.""" | ||
| async with Client(mcp) as client: | ||
| with pytest.raises(ToolError): | ||
| # Invalid input shape passes the protocol layer but fails handler-side | ||
| # validation before the wrapped handler runs, so this is protocol-level. | ||
| await client.call_tool("get_course_generation_prompt_tool", {"course_params": "not-a-mapping"}) | ||
|
|
||
| rows = await audit_events() | ||
| failures = [row for row in rows if row.action == "failed"] | ||
| assert len(failures) == 1 |
There was a problem hiding this comment.
Medium (test coverage): the real "handler raised → middleware stays silent" path isn't exercised.
As the comment on line 46 notes, this test triggers a protocol-level (pre-handler) validation failure — handler_recorded() is never set to True here, so it actually verifies the middleware records once, not that the middleware suppresses a duplicate.
The subtle part of the design — a wrapped handler that raises during execution sets _execution_recorded=True, and the middleware must then observe that and stay silent — depends on the contextvar mutation propagating from the handler back up through FastMCP's on_call_tool. That end-to-end path has no direct test. The success path (test_direct_server_tool_call_is_audited) exercises propagation for a returning handler, but not for a raising one.
Suggest registering a server tool whose handler raises at runtime (e.g. via a throwaway FastMCP instance + ToolCallAuditMiddleware), calling it, and asserting exactly one tool.invoked + one tool.failed (i.e. the backstop did not add a second failure). That pins the anti-double-record guarantee against future FastMCP dispatch changes.
There was a problem hiding this comment.
Added in 634b06f: test_runtime_handler_failure_is_recorded_exactly_once spins up a throwaway FastMCP server with the middleware, calls a tool that raises at runtime, and asserts exactly one tool.invoked + one tool.failed, i.e. the backstop stayed silent.
This comment was written with the assistance of an LLM (Claude).
| return await call_next(context) | ||
| # NOTE: Kept broad intentionally: any protocol-layer error | ||
| # (unknown tool, validation, handler crash) must be considered | ||
| # for the backstop record, and is always re-raised. | ||
| except Exception as exc: | ||
| if not handler_recorded(): | ||
| await self._record_protocol_failure(context.message, exc) | ||
| raise |
There was a problem hiding this comment.
When call_next fails because the wrapper's tool.invoked write raised, handler_recorded() is False, so this issues another record_event_now against the same (down) audit DB. That write also raises, replacing the original exception and leaving no record, which is the opposite of the module docstring's "every tools/call … leaves a record either way." Consider catching the backstop write's own failure (log it, preserve/chain the original exc) so an audit-store outage doesn't swap the surfaced error and lose the trail.
There was a problem hiding this comment.
Good catch, fixed in 634b06f. The backstop write is now wrapped in try/except SQLAlchemyError: we log it and re-raise the original error untouched. Also toned down the docstring, since when the audit DB itself is down there is genuinely nothing we can save. Test: test_backstop_write_failure_does_not_replace_the_original_error.
This comment was written with the assistance of an LLM (Claude).
| with ai_audit_context(source=AuditSource.CHAT, model=self._audit_model_info(response)): | ||
| tool_result = await self._execute_tool(tool_name, tool_args, tools) |
There was a problem hiding this comment.
Caught by Claude:
This PR installs ai_audit_context around tool execution to make capture fail-closed, but the wrapper's refusal exception is absorbed before it can fail the request:
- Chat:
_execute_tool'sbroadexcept Exceptionatproviders.py:324-326catches the wrapper'sSQLAlchemyErrorand returns it as a string;tool_funcatplugins/chat/tools.py:186only catches (ValidationError,ValueError,TypeError,RuntimeError), so the DB error flows straight to that broad handler. - RAG:
create_agent(...)atagent.py:65runs tools through LangGraph's ToolNode (defaulthandle_tool_errors=True), which converts the exception into an errorToolMessage;run_agentic_rag_retrievalonly re-raisesValidationError/GraphRecursionError/LangChainException.
Net: when the tool.invoked write fails, the call is refused (side effect prevented) but there's no audit record and no surfaced failure — it looks like an ordinary tool error and chat/RAG continue. That defeats "record on every execution path" for exactly these two seams. Suggest letting audit-infra exceptions propagate as a distinct hard failure past both seams, or documenting chat/RAG capture as best-effort.
There was a problem hiding this comment.
Fixed the chat half in 634b06f. The wrapper now raises a dedicated AuditCaptureError when the invoked/completed write fails, and _execute_tool re-raises it before its broad except, so an audit outage fails the request instead of becoming an error string for the model. Test: test_audit_capture_outage_is_a_hard_failure_not_a_tool_error_string.
The RAG half turns out not to apply on our pinned versions: langgraph 1.0.7's default handler only converts ToolInvocationError into a ToolMessage and re-raises everything else (handle_tool_errors=True was the old create_react_agent default). So the refusal already propagates there; documented it in the docstring.
This comment was written with the assistance of an LLM (Claude).
| except Exception as exc: | ||
| await record_event_now( | ||
| ToolFailedAuditEvent( |
There was a problem hiding this comment.
As this write is before the raise, so if it itself raises (e.g. if the DB is down), then the callers/logs will see an audit-infra error instead of the actual failure.
Perhaps guard the audit write as well:
except Exception as exc:
try:
await record_event_now(ToolFailedAuditEvent(..., error_detail=f"{type(exc).__name__}: {exc}"))
except SQLAlchemyError:
logger.exception("audit tool.failed write failed for %s", handler.__name__)
raise
There was a problem hiding this comment.
Fixed in 634b06f, pretty much as you suggested: the tool.failed write is guarded, logged, and the handler's own error re-raised. Since the invocation is already committed at that point, an invoked event with no outcome is exactly the crash signal, so we lose nothing. Test: test_failed_write_failure_surfaces_the_handler_error.
This comment was written with the assistance of an LLM (Claude).
| ) | ||
| _execution_recorded.set(True) | ||
| try: | ||
| result = await handler(*args, **kwargs) |
There was a problem hiding this comment.
await handler(*args, **kwargs) assumes a coroutine, but Tool.__post_init__ (lib/mcp/hooks.py:34-35) wraps every Tool.handler and both FastMCP and chat's sync_tool_func (plugins/chat/tools.py:190) support sync handlers. A sync handler would write tool.invoked, then hit TypeError: object … can't be used in 'await', recorded as tool.failed on every call. Latent today (all canvas/openedx handlers are async), but the wrapper narrows the accepted contract without a guard — consider inspect.iscoroutinefunction (fail loudly at wrap time) or supporting sync handlers.
Caught by Claude.
There was a problem hiding this comment.
Fixed in 634b06f, went with the fail-loudly option: TypeError at wrap time if the handler isn't a coroutine function. Supporting sync handlers would just move the problem (chat's sync_tool_func assumes a coroutine too). The plugins guide now says handlers must be async. Test: test_sync_handler_is_rejected_at_wrap_time.
This comment was written with the assistance of an LLM (Claude).
|
|
||
| @functools.wraps(handler) | ||
| async def audited(*args: Any, **kwargs: Any) -> Any: | ||
| call = AuditToolCall(name=handler.__name__, args=_named_jsonable_args(handler, args, kwargs)) |
There was a problem hiding this comment.
What if two different handlers have the same __name__ , but in different modules? Will they be tracked under the same name as well?
There was a problem hiding this comment.
Yes, they'd show up under the same tool_name; the row only records handler.__name__. In practice tool names are already the identity everywhere (the server refuses duplicate registrations and the LLM calls tools by name), so a second same-named tool can't actually be exposed. If you'd like provenance anyway, adding handler.__module__ as a column is cheap; happy to do it as a follow-up.
This comment was written with the assistance of an LLM (Claude).
| audited handler wrapper (:func:`sparkth.lib.audit.audited_tool_handler`); this | ||
| middleware covers the calls that never reach a handler (unknown tool names | ||
| and input-validation rejections), so every ``tools/call`` on ``/ai/mcp`` | ||
| leaves a record either way (ADR-0002 seam table row three). |
There was a problem hiding this comment.
I don't understand the point of tracking calls to tools that do not exist. This adds extra complexity without any visible gain. The way I see it, this is causing us to record things that did not happen. But I might be missing some context?
And if we do want to keep track of non-existing tool calls, then we should keep this middleware but get rid of the audited_tool_handler decorator. The decorator forces developers to wrap all tools, while the middleware automatically captures all calls. Isn't that a duplicated feature?
There was a problem hiding this comment.
A rejected tools/call is still a request that happened; we record it with outcome=failure, not as something that succeeded. Failed attempts (unknown tool, bad input, someone probing tool names) are exactly what an audit trail should have on record, same reason we record failed logins.
They also aren't duplicates, because they cover different ground: chat and RAG run tools in-process and never touch the middleware, so the wrapper is the only capture there; and calls that die before dispatch never reach a wrapped handler, so the middleware is the only capture for those. Developers don't wrap anything by hand either, Tool.__post_init__ does it automatically.
If we decide pre-dispatch failures aren't worth recording, the middleware (and the contextvar with it) can go, but that's a scope decision on the trail. Happy to discuss.
This comment was written with the assistance of an LLM (Claude).
| # Scoped per protocol-level call by tool_execution_recording(); the FastMCP | ||
| # middleware uses it to record failures that never reached a wrapped handler | ||
| # without ever double-recording ones that did. | ||
| _execution_recorded: ContextVar[bool] = ContextVar("audit_tool_execution_recorded", default=False) |
There was a problem hiding this comment.
This global context variable is a red flag to me. I believe that we need it because of the ToolCallAuditMiddleware, which tracks calls to tools that don't exist. If we get rid of ToolCallAuditMiddleware then we can get rid of this variable and the tool_execution_recording context manager, right?
There was a problem hiding this comment.
Right: it only exists so the wrapper and the middleware don't double-record on the MCP surface, so it lives and dies with the middleware. If the other thread lands on dropping the middleware, this and tool_execution_recording go with it. While we keep it, a contextvar is the async-safe way to share that one bit without threading state through FastMCP.
This comment was written with the assistance of an LLM (Claude).
| def _source(scope: Scope) -> AuditSource: | ||
| """MCP for the FastMCP mount (mounted at ``/ai`` in main.py), REST otherwise.""" | ||
| path: str = scope.get("path", "") | ||
| if path == "/ai" or path.startswith("/ai/"): |
There was a problem hiding this comment.
This "/ai" is a hard-coded string that should be replaced by a variable. Unless each router configures its own AuditSource?
There was a problem hiding this comment.
Fixed in 2982a4e: extracted MCP_MOUNT_PATH into core/config.py, used by both the mount in main.py and the middleware. Kept it a plain constant since we only have one non-REST mount today; easy to promote to per-router config if that changes.
This comment was written with the assistance of an LLM (Claude).
| _execution_recorded.reset(token) | ||
|
|
||
|
|
||
| def audited_tool_handler(handler: AsyncToolHandler) -> AsyncToolHandler: |
There was a problem hiding this comment.
I have a problem with the name of this decorator. If we decide to keep it, and not the middleware, then we should probably rename to something like "audited_tool".
In general "handler"/"manager"/etc. suffixes don't add meaningful to a name and should be either removed or replaced by something more specific. (There are exceptions though)
There was a problem hiding this comment.
Agreed, audited_tool is better. Holding the rename until the middleware question (other thread) settles, then I'll do it in one pass since it touches the lib surface and the docs.
This comment was written with the assistance of an LLM (Claude).
Review findings on PR #523: an audit-store outage could mask the real error on every path, or be silently absorbed. Invoked/completed write failures now surface as a distinct AuditCaptureError so the chat loop's broad tool-error handling can no longer turn a fail-closed refusal into a model-visible string; a failing tool.failed write (and the MCP backstop write) is logged and the original exception re-raised, since the committed invocation with no outcome event is already the abnormal-termination signal. error_detail is scrubbed before persistence (validation errors re-rendered without inputs, secret-keyed assignments masked, length bounded), and sync handlers are rejected at wrap time instead of failing on await at every call.
The "/ai" literal was duplicated between the app assembly and the audit context middleware's source detection (PR #523 review).
3 participants
Closes: AI tool executions are not audited on any execution path
What
Every AI tool execution, on every execution path (the FastMCP server on
/ai/mcp, the chat LangChain loop, and the RAG LangGraph agent), now leaves an append-only audit record with actor, origin, arguments, model identity, and outcome, closing the epic #435 gap for AI actions.Changes
audited_tool_handlerwrapper records a fail-closedtool.invokedevent before the handler runs (a failed write refuses the call, per ADR-0002 / NIST AU-5) plus atool.completedortool.failedoutcome event; both share a generated invocation id ontargetso a crash mid-execution is detectabletool.invoked/tool.completed/tool.failedevent types on theAUDIT_EVENTShook, andai_audit_contextto stamp the surface (chat/rag) and driving model onto tool executionsToolhook dataclass audit-wraps its handler at construction, covering the FastMCP server and the chat tool registry with one seam; wrapping preserves name, docstring, and generated input schema (regression-tested)chatTool.handler) wrap explicitly; the agent runner stamps sourceragwith best-effort model identityToolCallAuditMiddlewarebackstop records protocol-level failures that never reach a handler (unknown tool, input validation) without double-recording handler failures; the audit context middleware stamps sourcemcpfor the/aimountHow to Test
make services.up && make migrations && make backend.up.devhttp://localhost:7727/ai/mcp:call_tool("get_course_generation_prompt_tool", {"course_params": {"course_name": "Algebra", "course_description": "Linear equations"}})make db-shell, thenSELECT category, action, source, outcome, tool_name, target_id, tool_args FROM audit_events ORDER BY id;shows thetool.invoked/tool.completedpair withsource=mcp, a sharedtarget_id, and captured request IP / request idtool.failedrow appears (protocol-level backstop)canvas_authenticatewith a bogus URL and token: theauthpayload is stored as"[REDACTED]"and the failed execution recordstool.invoked+tool.failedsource=chatplusmodel_provider/model_name/model_version(covered bytests/llm/test_provider_tool_audit.py)make test.backend.pytest && make mypy && make lint.backendall pass (1439 tests)Notes
No migration: the
audit_eventsAI-provenance columns shipped with #492. Fail-closed semantics mean an audit-database outage refuses AI tool calls instead of running them unrecorded; this is the deliberate ADR-0002 trade-off. Argument-validation refusals inside the chat tool conversion layer (before a handler is reached) are not yet recorded; auxiliary LLM calls without tools (title generation, scope classifier, Slack synthesis) are generation-level provenance and remain for a follow-up issue.This pull request (code and description) was written with the assistance of an LLM (Claude).