Skip to content

Commit e779def

Browse files
Dumbrisclaude
andauthored
fix(server): close the tool_call record when output_sanitisation blocks a response (#1687)
## Summary When `output_sanitisation` blocks a response (`response_action: block`), the call_tool_* handlers returned before recording the call's completion. The upstream call **did** run (side effects may have happened), yet only a `policy_decision` was persisted: no `tool_call` activity row (started events are not persisted), an SSE `started` event with no completion, and no tool-call history entry. `handleCallToolVariant`, the legacy `handleCallTool` and direct routing mode (`makeDirectModeHandler`, added by #1685) now close the call on a block: - **Tool-call history:** the attempt is stored with `Error` set to the block message and `Response` cleared, so the blocked payload is never stored. - **Activity:** a `tool_call` completion is emitted with status `error`, `error_class: proxy_policy` and `fault_domain: proxy`. It carries the fixed block message, never the response. On call_tool_* paths the `internal_tool_call` record is emitted as well. - The new helper `recordOutputBlocked` follows the convention of the code_execution sub-call block path, which already closed its record this way. ## Counting / audit check - **Spec 107 audit:** the block's `policy_decision` already writes `tool_call outcome:blocked reason:output_sanitisation` (`isPostDispatchBlockKey`). The completion's audit write is the existing deduped no-op (`toolCallWritten`), so the audit log still has one line per attempt. - **Telemetry:** nothing is double-counted. The availability block counter comes only from `policy_decision`, the upstream-attempt bucket was counted at dispatch, and activation is not recorded. - **Activity/Usage call totals (`storage.CountsAsCall`):** the `policy_decision` (blocked) and the new `tool_call` (error) both count as failed calls. One blocked attempt therefore counts as two failed calls in those totals. The code_execution sanitisation block and the generation-changed refusal pair already counted this way. Separating them needs the post-dispatch block key on the `policy_decision` record, so that is left for a follow-up. ## Direct mode After #1685 merged, origin/main was merged into this branch and the direct-mode block now emits the same error `tool_call` completion. Direct mode keeps no tool-call history. ## Tests - `TestCallToolVariant_SanitisationBlockRecordsToolCallCompletion` and `TestLegacyCallTool_SanitisationBlockRecordsToolCallCompletion` drive a real in-process upstream that returns an AWS key with block mode on. Each asserts one persisted `tool_call` with status `error` and `error_class` `proxy_policy`, the `policy_decision` beside it, and the secret absent from every activity record. `TestDirectMode_SanitisationBlockRecordsToolCallCompletion` asserts the same for direct mode. All three failed before the fix (no `tool_call` record). - `TestRecordOutputBlocked_HistoryCarriesBlockMessageNotResponse` asserts the history entry carries the block message and no response. The direct-routing fixture's incarnation guard drops all history writes, so this is tested on the helper. - `go test -race ./internal/server/` with the unit-tests.yml `-skip` regex: pass. - `golangci-lint` v2 (`.github/.golangci.yml`), bare and `--build-tags server`, over `./internal/server/...`: no issues in changed files. Six existing staticcheck findings remain in untouched test files. 🤖 Generated with [Claude Code](https://claude.com/claude-code) --------- Co-authored-by: Claude Opus 5.5 <noreply@anthropic.com>
1 parent ab0574c commit e779def

3 files changed

Lines changed: 179 additions & 0 deletions

File tree

‎internal/server/mcp.go‎

Lines changed: 31 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -3564,6 +3564,16 @@ func (p *MCPProxyServer) handleCallToolVariant(ctx context.Context, request mcp.
35643564
// holds an unredacted secret and a blocked response is never cached. A
35653565
// non-nil result means the call was blocked.
35663566
if blockResult := p.applyOutputSanitisation(ctx, serverName, actualToolName, requestID, contentTrust, result); blockResult != nil {
3567+
// The upstream ran, so the started record is closed: history and the
3568+
// tool_call completion carry the block message, never the response.
3569+
blockMsg := p.recordOutputBlocked(ctx, toolCallRecord, callInc, blockResult)
3570+
var intentMap map[string]interface{}
3571+
if intent != nil {
3572+
intentMap = intent.ToMap()
3573+
}
3574+
p.emitActivityToolCallCompleted(ctx, serverName, actualToolName, sessionID, requestID, activitySource, storage.ActivityStatusError, blockMsg, duration.Milliseconds(), activityArgs, "", false, toolVariant, intentMap, contentTrust, profileSlug, 0, 0, "", nil, "")
3575+
internalToolName := "call_tool_" + intent.OperationType
3576+
p.emitActivityInternalToolCall(ctx, internalToolName, serverName, actualToolName, toolVariant, sessionID, requestID, storage.ActivityStatusError, blockMsg, time.Since(internalStartTime).Milliseconds(), activityArgs, nil, intentMap, "")
35673577
return blockResult, nil
35683578
}
35693579
// Spec 112: sanitisation mutates result in place, so the scrubbed record
@@ -4054,6 +4064,9 @@ func (p *MCPProxyServer) handleCallTool(ctx context.Context, request mcp.CallToo
40544064
// holds an unredacted secret and a blocked response is never cached. A
40554065
// non-nil result means the call was blocked.
40564066
if blockResult := p.applyOutputSanitisation(ctx, serverName, actualToolName, requestID, contentTrust, result); blockResult != nil {
4067+
// The upstream ran: close the started record (see handleCallToolVariant).
4068+
blockMsg := p.recordOutputBlocked(ctx, toolCallRecord, callInc, blockResult)
4069+
p.emitActivityToolCallCompleted(ctx, serverName, actualToolName, sessionID, requestID, activitySource, storage.ActivityStatusError, blockMsg, duration.Milliseconds(), activityArgs, "", false, "", nil, "", "", 0, 0, "", nil, "")
40574070
return blockResult, nil
40584071
}
40594072
// Spec 112: sanitisation mutates result in place, so the scrubbed record
@@ -4129,6 +4142,24 @@ func (p *MCPProxyServer) toolCallIncarnation(serverName string) int64 {
41294142

41304143
// recordToolCallAt records a call's history only while its server still has
41314144
// the incarnation the call started on (Spec 116 FR-014c).
4145+
// recordOutputBlocked is the bookkeeping a dispatch owes a call whose result
4146+
// an output policy (Spec 054 block mode) refused to hand over: the upstream
4147+
// ran, so the tool-call history gets the attempt, with the block message in
4148+
// place of the response it must never store. It notes the outcome as a proxy
4149+
// policy refusal (the upstream answered fine) and returns the message for the
4150+
// caller's tool_call completion, which closes the started record. The audit
4151+
// tool_call line is already written by the block's policy_decision.
4152+
func (p *MCPProxyServer) recordOutputBlocked(ctx context.Context, rec *storage.ToolCallRecord, inc int64, blockResult *mcp.CallToolResult) string {
4153+
msg := firstTextOf(blockResult)
4154+
rec.Response = nil
4155+
rec.Error = msg
4156+
if err := p.recordToolCallAt(rec, inc); err != nil {
4157+
p.logger.Warn("Failed to record blocked tool call", zap.Error(err))
4158+
}
4159+
callerr.NoteOutcome(ctx, callerr.Outcome{Class: callerr.ClassProxyPolicy, Domain: callerr.DomainProxy})
4160+
return msg
4161+
}
4162+
41324163
func (p *MCPProxyServer) recordToolCallAt(rec *storage.ToolCallRecord, inc int64) error {
41334164
written, err := p.storage.RecordToolCallIfIncarnation(rec, inc)
41344165
if err == nil && !written {

‎internal/server/mcp_routing.go‎

Lines changed: 5 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -745,6 +745,11 @@ func (p *MCPProxyServer) makeDirectModeHandler(entry *directCatalogEntry) mcpser
745745
// applied to /mcp/all. Applied to the RAW result before the activity
746746
// classification and truncation below; a non-nil result means blocked.
747747
if blockResult := p.applyOutputSanitisation(ctx, serverName, toolName, requestID, directContentTrust, result); blockResult != nil {
748+
// The upstream ran, so the started record is closed with the block
749+
// message, never the response (direct mode keeps no history).
750+
blockMsg := firstTextOf(blockResult)
751+
callerr.NoteOutcome(ctx, callerr.Outcome{Class: callerr.ClassProxyPolicy, Domain: callerr.DomainProxy})
752+
p.emitActivityToolCallCompleted(ctx, serverName, toolName, sessionID, requestID, "mcp", "error", blockMsg, durationMs, enrichedArgs, "", false, contracts.DeriveCallWith(annotations), nil, directContentTrust, "", 0, 0, "", nil, "")
748753
return blockResult, nil
749754
}
750755

Lines changed: 143 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -0,0 +1,143 @@
1+
package server
2+
3+
import (
4+
"context"
5+
"encoding/json"
6+
"testing"
7+
"time"
8+
9+
"github.com/mark3labs/mcp-go/mcp"
10+
"github.com/stretchr/testify/assert"
11+
"github.com/stretchr/testify/require"
12+
13+
"github.com/smart-mcp-proxy/mcpproxy-go/internal/config"
14+
"github.com/smart-mcp-proxy/mcpproxy-go/internal/contracts"
15+
"github.com/smart-mcp-proxy/mcpproxy-go/internal/security"
16+
"github.com/smart-mcp-proxy/mcpproxy-go/internal/storage"
17+
)
18+
19+
// When output_sanitisation blocks a response the upstream call DID run, so
20+
// the dispatch must still close its started record with a tool_call
21+
// completion: without one the Activity Log has no row for a call that may
22+
// have had side effects, and SSE consumers see a start that never ends. The
23+
// completion carries the block message, never the response that was blocked.
24+
25+
// blockedSecretFixture is a direct-routing fixture whose "github:get_secret"
26+
// answers with a critical secret, block mode on, and the activity service
27+
// persisting records.
28+
func blockedSecretFixture(t *testing.T) *directRoutingFixture {
29+
t.Helper()
30+
f := newDirectRoutingFixture(t, directServerSpec{server: "github", tools: []toolSpec{readSpec("get_secret")}})
31+
f.ups["github"].mcpSrv.AddTool(mcp.Tool{Name: "get_secret", InputSchema: mcp.ToolInputSchema{Type: "object"}},
32+
func(_ context.Context, req mcp.CallToolRequest) (*mcp.CallToolResult, error) {
33+
f.ups["github"].record(req.Params.Name)
34+
return mcp.NewToolResultText("leaked " + awsKeyFixture), nil
35+
})
36+
f.proxy.config.OutputSanitisation = sanCfg("block", false, false)
37+
f.proxy.sanitisationDetector = security.NewDetector(config.DefaultSensitiveDataDetectionConfig())
38+
39+
go f.rt.ActivityService().Start(f.rt.AppContext(), f.rt)
40+
deadline := time.Now().Add(5 * time.Second)
41+
for !f.rt.ActivityService().Started() && time.Now().Before(deadline) {
42+
time.Sleep(time.Millisecond)
43+
}
44+
require.True(t, f.rt.ActivityService().Started(), "activity service must subscribe before the call")
45+
return f
46+
}
47+
48+
// assertBlockedCallRecorded checks the persisted activity and tool-call
49+
// history for the blocked call: one tool_call row in error state with the
50+
// block message, the policy_decision beside it, and the secret nowhere.
51+
func assertBlockedCallRecorded(t *testing.T, f *directRoutingFixture) {
52+
t.Helper()
53+
rec := awaitToolCallActivity(t, f.rt.StorageManager(), "github", "get_secret")
54+
assert.Equal(t, storage.ActivityStatusError, rec.Status)
55+
assert.Contains(t, rec.ErrorMessage, "blocked by sanitisation policy")
56+
assert.Equal(t, "proxy_policy", rec.ErrorClass, "an output-policy block is mcpproxy's refusal, not an upstream failure")
57+
raw, err := json.Marshal(rec)
58+
require.NoError(t, err)
59+
assert.NotContains(t, string(raw), awsKeyFixture, "the tool_call record must not carry the blocked response")
60+
61+
all, _, err := f.rt.StorageManager().ListActivities(storage.ActivityFilter{Limit: 100})
62+
require.NoError(t, err)
63+
toolCalls, blocks := 0, 0
64+
for _, r := range all {
65+
if r.ServerName != "github" || r.ToolName != "get_secret" {
66+
continue
67+
}
68+
switch r.Type {
69+
case storage.ActivityTypeToolCall:
70+
toolCalls++
71+
case storage.ActivityTypePolicyDecision:
72+
if r.Status == storage.ActivityStatusBlocked {
73+
blocks++
74+
}
75+
}
76+
b, mErr := json.Marshal(r)
77+
require.NoError(t, mErr)
78+
assert.NotContains(t, string(b), awsKeyFixture, "no activity record may carry the blocked response")
79+
}
80+
assert.Equal(t, 1, toolCalls, "exactly one tool_call completion")
81+
assert.Equal(t, 1, blocks, "the block decision is still recorded")
82+
83+
}
84+
85+
func TestCallToolVariant_SanitisationBlockRecordsToolCallCompletion(t *testing.T) {
86+
f := blockedSecretFixture(t)
87+
req := mcp.CallToolRequest{}
88+
req.Params.Name = contracts.ToolVariantRead
89+
req.Params.Arguments = map[string]interface{}{"name": "github:get_secret", "args": map[string]interface{}{}}
90+
res, err := f.proxy.handleCallToolVariant(fullTierAgentOn("github"), req, contracts.ToolVariantRead)
91+
require.NoError(t, err)
92+
require.True(t, res.IsError)
93+
require.Equal(t, int64(1), f.ups["github"].count.Load(), "the upstream ran")
94+
95+
assertBlockedCallRecorded(t, f)
96+
}
97+
98+
func TestLegacyCallTool_SanitisationBlockRecordsToolCallCompletion(t *testing.T) {
99+
f := blockedSecretFixture(t)
100+
req := mcp.CallToolRequest{}
101+
req.Params.Name = "call_tool"
102+
req.Params.Arguments = map[string]interface{}{"name": "github:get_secret", "args": map[string]interface{}{}}
103+
res, err := f.proxy.handleCallTool(adminCtx(), req)
104+
require.NoError(t, err)
105+
require.True(t, res.IsError)
106+
require.Equal(t, int64(1), f.ups["github"].count.Load(), "the upstream ran")
107+
108+
assertBlockedCallRecorded(t, f)
109+
}
110+
111+
func TestDirectMode_SanitisationBlockRecordsToolCallCompletion(t *testing.T) {
112+
f := blockedSecretFixture(t)
113+
res := f.call(t, fullTierAgentOn("github"), "github", "get_secret")
114+
require.True(t, res.IsError)
115+
require.Equal(t, int64(1), f.ups["github"].count.Load(), "the upstream ran")
116+
117+
assertBlockedCallRecorded(t, f)
118+
}
119+
120+
// The tool-call history entry of a blocked call keeps the attempt but never
121+
// the response: the block message takes its place. (The direct-routing
122+
// fixture's incarnation guard drops every history write, so this is checked
123+
// on the helper the dispatch paths share.)
124+
func TestRecordOutputBlocked_HistoryCarriesBlockMessageNotResponse(t *testing.T) {
125+
proxy, _ := createTestProxyWithRuntime(t, nil)
126+
rec := &storage.ToolCallRecord{
127+
ID: "call-1", ServerID: "srv-id", ServerName: "github", ToolName: "get_secret",
128+
Response: mcp.NewToolResultText("leaked " + awsKeyFixture),
129+
Timestamp: time.Now(),
130+
}
131+
msg := proxy.recordOutputBlocked(context.Background(), rec, 0,
132+
mcp.NewToolResultError("tool output blocked by sanitisation policy: critical sensitive data detected in tool output"))
133+
assert.Contains(t, msg, "blocked by sanitisation policy")
134+
135+
history, err := proxy.storage.GetServerToolCalls("srv-id", 10)
136+
require.NoError(t, err)
137+
require.Len(t, history, 1)
138+
assert.Equal(t, msg, history[0].Error)
139+
assert.Nil(t, history[0].Response)
140+
raw, err := json.Marshal(history[0])
141+
require.NoError(t, err)
142+
assert.NotContains(t, string(raw), awsKeyFixture)
143+
}

0 commit comments

Comments
 (0)