Skip to content

Commit c3ed54f

Browse files
authored
fix(082): stamp the work session on every tool call, not only the ones that remembered to (#843)
One opencode connection showed up in the Activity Log session picker as TWO sessions, 20 seconds apart. Both halves were the same connection: 05:30:39 list_registries work_session_id = null 05:30:41 upstream_servers work_session_id = null 05:30:45 retrieve_tools work_session_id = ws-a3f505743e000abe 05:30:46 echo work_session_id = ws-a3f505743e000abe The labels in the picker were the tails of two DIFFERENT id namespaces — `58cea` from the transport session id, `00abe` from the work session id. Backend: work was marked inside individual tool handlers (retrieve_tools, the call_tool_* variants, code_execution). Every other built-in — list_registries, upstream_servers, quarantine_security — wrote activity without ever resolving a work session. Marking now happens in the beforeAny hook, which every method passes through on every server instance (direct / code-exec / call-tool share the hooks), so a new built-in cannot reintroduce this by forgetting a call. tools/call remains the boundary of "work": a connection that only initializes and lists tools still persists nothing, which is what keeps background agents from burying the user's real sessions. Frontend: grouping by `work_session_id || session_id` is what turned an orphaned row into a phantom session. It now resolves the fallback through an index (transport session -> work session, learned from the sessions API and from any sibling row that carries one), so rows with no work session of their own fold into their connection's instead of keying on the raw transport id. This also heals the rows already in the database. Verified against a real MCP session on a scratch instance: a connection that calls upstream_servers and then retrieve_tools now produces two records sharing one work_session_id, and the picker shows one entry, not two. Spec 082 SC-002.
1 parent c603654 commit c3ed54f

6 files changed

Lines changed: 306 additions & 10 deletions

File tree

Lines changed: 49 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -0,0 +1,49 @@
1+
import type { ActivityRecord, MCPSession } from '@/types/api'
2+
3+
/**
4+
* Spec 082 — which session an activity row belongs to.
5+
*
6+
* The row we want to group by is the WORK session: one client, one project,
7+
* across reconnects. But not every row carries one. A record written before the
8+
* connection was attributed — or by an older build, or by a path that never
9+
* marked the session as worked — has only its TRANSPORT session id.
10+
*
11+
* Falling back to the transport id per-row (`work_session_id || session_id`) is
12+
* what produced the original bug: the orphaned rows of a connection landed in a
13+
* bucket keyed by the transport id while the rest of that same connection landed
14+
* in its work-session bucket, and one client showed up in the picker twice.
15+
*
16+
* So resolve the fallback through an index instead: a transport session belongs
17+
* to whatever work session ANY of its rows (or the sessions API) names. Only a
18+
* transport session with no known work session anywhere keys on itself.
19+
*/
20+
export type WorkSessionIndex = Map<string, string>
21+
22+
export function buildWorkSessionIndex(
23+
activities: readonly ActivityRecord[],
24+
sessions: readonly MCPSession[] = [],
25+
): WorkSessionIndex {
26+
const index: WorkSessionIndex = new Map()
27+
28+
// The sessions API is the more authoritative of the two: it is the connection's
29+
// own record of the work session it was attributed to.
30+
for (const s of sessions) {
31+
if (s.id && s.work_session_id) index.set(s.id, s.work_session_id)
32+
}
33+
34+
// Sibling rows fill the gaps — activity outlives sessions (90 days vs the 100
35+
// most recent), so for older rows this is the only surviving evidence.
36+
for (const a of activities) {
37+
if (a.session_id && a.work_session_id && !index.has(a.session_id)) {
38+
index.set(a.session_id, a.work_session_id)
39+
}
40+
}
41+
42+
return index
43+
}
44+
45+
export function groupKeyOf(a: ActivityRecord, index: WorkSessionIndex): string {
46+
if (a.work_session_id) return a.work_session_id
47+
if (a.session_id) return index.get(a.session_id) ?? a.session_id
48+
return ''
49+
}

frontend/src/views/Activity.vue

Lines changed: 18 additions & 5 deletions
Original file line numberDiff line numberDiff line change
@@ -693,8 +693,9 @@ import { ref, computed, onMounted, onUnmounted, watch } from 'vue'
693693
import { useRoute } from 'vue-router'
694694
import { useSystemStore } from '@/stores/system'
695695
import api from '@/services/api'
696-
import type { ActivityRecord, ActivitySummaryResponse } from '@/types/api'
696+
import type { ActivityRecord, ActivitySummaryResponse, MCPSession } from '@/types/api'
697697
import { buildSessionLabels } from '@/utils/sessionLabel'
698+
import { buildWorkSessionIndex, groupKeyOf as workSessionKeyOf } from '@/utils/sessionGrouping'
698699
import JsonViewer from '@/components/JsonViewer.vue'
699700
700701
const route = useRoute()
@@ -776,6 +777,10 @@ interface SessionOption {
776777
// show "Claude Code · 14:32" instead of an opaque "...139c9".
777778
const sessionInfo = ref(new Map<string, { clientName?: string; startTime?: string; workspace?: string }>())
778779
780+
// The raw session records behind that map. Kept because they carry the
781+
// transport -> work session link that groupKeyOf needs (see workSessionIndex).
782+
const sessionsRaw = ref<MCPSession[]>([])
783+
779784
// Session ids we have already tried and failed to resolve. The core keeps only
780785
// the 100 most recent sessions, so an old session's activity rows can outlive
781786
// its record and never become resolvable. Without this set, every refresh would
@@ -791,8 +796,10 @@ const loadSessions = async () => {
791796
sessionsInFlight = (async () => {
792797
try {
793798
const response = await api.getSessions(100)
799+
const sessions = response.data?.sessions ?? []
800+
sessionsRaw.value = sessions
794801
const next = new Map<string, { clientName?: string; startTime?: string; workspace?: string }>()
795-
for (const s of response.data?.sessions ?? []) {
802+
for (const s of sessions) {
796803
const info = {
797804
clientName: s.client_name,
798805
startTime: s.start_time,
@@ -844,10 +851,16 @@ const refreshSessionsIfUnknown = () => {
844851
if (hasUnknown) void loadSessions()
845852
}
846853
854+
// Transport session -> work session, learned from the sessions API and from any
855+
// sibling row that does carry one. A row with no work session of its own is
856+
// folded into its connection's, rather than keying on the raw transport id and
857+
// splitting one client into two entries in the picker.
858+
const workSessionIndex = computed(() => buildWorkSessionIndex(activities.value, sessionsRaw.value))
859+
847860
// The key an activity row is grouped and filtered by: its WORK session (Spec
848-
// 082) — one client, one project, across reconnects. Rows written before 082
849-
// have none, so they fall back to the transport session and behave as before.
850-
const groupKeyOf = (a: ActivityRecord): string => a.work_session_id || a.session_id || ''
861+
// 082) — one client, one project, across reconnects. Rows for which no work
862+
// session is known anywhere still fall back to the transport session.
863+
const groupKeyOf = (a: ActivityRecord): string => workSessionKeyOf(a, workSessionIndex.value)
851864
852865
const availableSessions = computed((): SessionOption[] => {
853866
const seen = new Map<string, { clientName?: string; startTime?: string; workspace?: string }>()
Lines changed: 80 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -0,0 +1,80 @@
1+
import { describe, it, expect } from 'vitest'
2+
import { buildWorkSessionIndex, groupKeyOf } from '@/utils/sessionGrouping'
3+
import type { ActivityRecord, MCPSession } from '@/types/api'
4+
5+
const rec = (p: Partial<ActivityRecord>): ActivityRecord =>
6+
({ status: 'success', timestamp: '2026-07-13T05:30:00Z', ...p }) as ActivityRecord
7+
8+
// The bug this file exists for (Spec 082 SC-002).
9+
//
10+
// One opencode connection, transport session `mcp-session-…858cea`, wrote four
11+
// activity records. The first two named built-ins that did not mark the session
12+
// as worked, so they carry no work_session_id; the last two carry `ws-…000abe`.
13+
// Grouping by `work_session_id || session_id` put them in two different buckets
14+
// and the session picker showed ONE connection as TWO sessions — labelled with
15+
// the tails of two different ids, `58cea` and `00abe`.
16+
describe('work-session grouping', () => {
17+
const SID = 'mcp-session-9958d45e-015f-4c16-8d9a-74792c858cea'
18+
const WSID = 'ws-a3f505743e000abe'
19+
20+
const opencodeRun: ActivityRecord[] = [
21+
rec({ session_id: SID, tool_name: 'list_registries' }), // no work_session_id
22+
rec({ session_id: SID, tool_name: 'upstream_servers' }), // no work_session_id
23+
rec({ session_id: SID, work_session_id: WSID, tool_name: 'retrieve_tools' }),
24+
rec({ session_id: SID, work_session_id: WSID, tool_name: 'echo' }),
25+
]
26+
27+
it('folds records with no work session into their connection’s work session', () => {
28+
const index = buildWorkSessionIndex(opencodeRun, [])
29+
const keys = new Set(opencodeRun.map(a => groupKeyOf(a, index)))
30+
31+
expect(keys).toEqual(new Set([WSID]))
32+
})
33+
34+
it('learns the mapping from the sessions API too, not only from sibling records', () => {
35+
// The orphaned rows on their own — the work session is only known because
36+
// /api/v1/sessions reports it for this transport session.
37+
const orphans = opencodeRun.slice(0, 2)
38+
const sessions = [{ id: SID, work_session_id: WSID } as MCPSession]
39+
40+
const index = buildWorkSessionIndex(orphans, sessions)
41+
42+
expect(orphans.map(a => groupKeyOf(a, index))).toEqual([WSID, WSID])
43+
})
44+
45+
it('falls back to the transport session when no work session is known anywhere', () => {
46+
// Pre-082 rows, and clients whose work session was never resolved.
47+
const legacy = [rec({ session_id: 'sess-legacy', tool_name: 'echo' })]
48+
const index = buildWorkSessionIndex(legacy, [])
49+
50+
expect(groupKeyOf(legacy[0], index)).toBe('sess-legacy')
51+
})
52+
53+
it('keeps genuinely different work sessions apart', () => {
54+
// Same client reconnecting into DIFFERENT work (e.g. a different project):
55+
// two transport sessions, two work sessions. These must not be merged.
56+
const rows = [
57+
rec({ session_id: 'sess-a', work_session_id: 'ws-project-one' }),
58+
rec({ session_id: 'sess-b', work_session_id: 'ws-project-two' }),
59+
]
60+
const index = buildWorkSessionIndex(rows, [])
61+
62+
expect(rows.map(a => groupKeyOf(a, index))).toEqual(['ws-project-one', 'ws-project-two'])
63+
})
64+
65+
it('groups a work session that spans several connections (the reconnect case)', () => {
66+
const rows = [
67+
rec({ session_id: 'sess-1', work_session_id: 'ws-same' }),
68+
rec({ session_id: 'sess-2', work_session_id: 'ws-same' }),
69+
rec({ session_id: 'sess-2' }), // orphan on the second connection
70+
]
71+
const index = buildWorkSessionIndex(rows, [])
72+
73+
expect(new Set(rows.map(a => groupKeyOf(a, index)))).toEqual(new Set(['ws-same']))
74+
})
75+
76+
it('tolerates records with no session id at all', () => {
77+
const index = buildWorkSessionIndex([rec({ tool_name: 'echo' })], [])
78+
expect(groupKeyOf(rec({ tool_name: 'echo' }), index)).toBe('')
79+
})
80+
})

internal/server/mcp.go

Lines changed: 22 additions & 3 deletions
Original file line numberDiff line numberDiff line change
@@ -27,6 +27,7 @@ import (
2727
"github.com/smart-mcp-proxy/mcpproxy-go/internal/outputvalidation"
2828
"github.com/smart-mcp-proxy/mcpproxy-go/internal/registries"
2929
"github.com/smart-mcp-proxy/mcpproxy-go/internal/reqcontext"
30+
"github.com/smart-mcp-proxy/mcpproxy-go/internal/runtime"
3031
"github.com/smart-mcp-proxy/mcpproxy-go/internal/security"
3132
"github.com/smart-mcp-proxy/mcpproxy-go/internal/server/tokens"
3233
"github.com/smart-mcp-proxy/mcpproxy-go/internal/shellwrap"
@@ -118,6 +119,11 @@ type MCPProxyServer struct {
118119
mainServer *Server // Reference to main server for config persistence
119120
config *config.Config // Add config reference for security checks
120121

122+
// workSessionResolver maps an identity to a work session (Spec 082). Normally
123+
// nil, and the runtime's tracker is used; tests inject a stub so they do not
124+
// have to stand up a whole Runtime to exercise session attribution.
125+
workSessionResolver func(runtime.WorkSessionIdentity) string
126+
121127
// Routing mode MCP server instances (Spec 031)
122128
// Each instance has different tools registered for its routing mode.
123129
directServer *mcpserver.MCPServer // Direct mode: upstream tools with serverName__toolName naming
@@ -235,10 +241,12 @@ func NewMCPProxyServer(
235241
})
236242
}
237243

238-
// Forward handle to the MCP server, which is constructed below but is needed
239-
// by the hooks above it (to send a roots request back to the client).
240-
// Atomic because the hooks run on client goroutines.
244+
// Forward handles to the MCP server and the proxy, both constructed below but
245+
// needed by the hooks above them (to send a roots request back to the client,
246+
// and to attribute work to a session). Atomic because the hooks run on client
247+
// goroutines.
241248
var mcpSrvRef atomic.Pointer[mcpserver.MCPServer]
249+
var proxyRef atomic.Pointer[MCPProxyServer]
242250

243251
// Create hooks to capture session information
244252
hooks := &mcpserver.Hooks{}
@@ -270,6 +278,14 @@ func NewMCPProxyServer(
270278
if method != mcp.MethodInitialize && sessionStore.TryClaimWorkspaceFetch(session.SessionID()) {
271279
go fetchWorkspaceRoot(ctx, mcpSrvRef.Load(), sessionStore, logger)
272280
}
281+
282+
// Spec 082: attribute this request to a work session before the handler
283+
// runs, so every activity record the handler writes carries it. Kicked off
284+
// AFTER the roots fetch above, which it may wait on briefly — the fetch is
285+
// already in flight by then.
286+
if proxy := proxyRef.Load(); proxy != nil {
287+
proxy.markWorkIfToolCall(ctx, method)
288+
}
273289
})
274290

275291
hooks.AddOnRegisterSession(func(ctx context.Context, sess mcpserver.ClientSession) {
@@ -432,6 +448,9 @@ func NewMCPProxyServer(
432448
hooks: hooks,
433449
}
434450

451+
// Let the hooks (registered before the proxy existed) reach it.
452+
proxyRef.Store(proxy)
453+
435454
// Register proxy tools for the default (retrieve_tools) server
436455
proxy.registerTools(debugSearch)
437456

Lines changed: 94 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -0,0 +1,94 @@
1+
package server
2+
3+
import (
4+
"context"
5+
"testing"
6+
7+
"github.com/mark3labs/mcp-go/mcp"
8+
mcpserver "github.com/mark3labs/mcp-go/server"
9+
"github.com/stretchr/testify/assert"
10+
"github.com/stretchr/testify/require"
11+
"go.uber.org/zap"
12+
13+
"github.com/smart-mcp-proxy/mcpproxy-go/internal/runtime"
14+
"github.com/smart-mcp-proxy/mcpproxy-go/internal/storage"
15+
)
16+
17+
// newWorkSessionProxy builds the smallest MCPProxyServer that can resolve a work
18+
// session: a real session store (persistence needs a storage manager) and a stub
19+
// resolver, so the test does not have to stand up a whole Runtime.
20+
func newWorkSessionProxy(t *testing.T) (*MCPProxyServer, *SessionStore) {
21+
t.Helper()
22+
23+
logger := zap.NewNop()
24+
sm, err := storage.NewManager(t.TempDir(), logger.Sugar())
25+
require.NoError(t, err)
26+
t.Cleanup(func() { _ = sm.Close() })
27+
28+
store := NewSessionStore(logger)
29+
store.SetStorageManager(sm)
30+
31+
p := &MCPProxyServer{
32+
logger: logger,
33+
sessionStore: store,
34+
workSessionResolver: func(id runtime.WorkSessionIdentity) string {
35+
return "ws-" + id.ClientName
36+
},
37+
}
38+
return p, store
39+
}
40+
41+
func workSessionCtx(sessionID string) context.Context {
42+
helper := mcpserver.NewMCPServer("test", "1.0.0")
43+
return helper.WithContext(context.Background(), &fakeClientSession{id: sessionID})
44+
}
45+
46+
// The regression: EVERY tools/call is work, whichever tool it names.
47+
//
48+
// Work used to be stamped inside individual handlers (retrieve_tools, the
49+
// call_tool_* variants, code_execution). The other built-ins — list_registries,
50+
// upstream_servers, quarantine_security, … — never marked the session, so the
51+
// activity they wrote carried an empty work_session_id. The Web UI then grouped
52+
// those rows under the raw TRANSPORT session id while the rest of the same
53+
// connection grouped under its work-session id, and one opencode connection
54+
// appeared in the session picker as two sessions (Spec 082 SC-002).
55+
func TestWorkSession_AnyToolCallMarksWork(t *testing.T) {
56+
p, store := newWorkSessionProxy(t)
57+
store.SetSession("s1", "opencode", "1.17.18", false, false, nil)
58+
59+
// A tools/call naming a built-in that never called markSessionWorked itself.
60+
got := p.markWorkIfToolCall(workSessionCtx("s1"), mcp.MethodToolsCall)
61+
62+
assert.Equal(t, "ws-opencode", got,
63+
"a tools/call must resolve a work session regardless of which tool it names")
64+
assert.Equal(t, "ws-opencode", store.WorkSessionID("s1"),
65+
"the work session must be cached on the connection, so every later record agrees")
66+
}
67+
68+
// The counterweight, and the reason this is a method allow-list rather than
69+
// "mark on any request": a connection that only shakes hands and lists tools has
70+
// done no work and must leave nothing behind. On a real machine 99 of 100
71+
// session records were background agents doing exactly this, every ~15 minutes.
72+
func TestWorkSession_HandshakeAndListingAreNotWork(t *testing.T) {
73+
p, store := newWorkSessionProxy(t)
74+
store.SetSession("s1", "background-agent", "1.0", false, false, nil)
75+
76+
for _, method := range []mcp.MCPMethod{
77+
mcp.MethodInitialize,
78+
mcp.MethodToolsList,
79+
mcp.MethodPing,
80+
} {
81+
assert.Empty(t, p.markWorkIfToolCall(workSessionCtx("s1"), method),
82+
"%s is not work and must not persist a session", method)
83+
}
84+
85+
assert.Empty(t, store.WorkSessionID("s1"),
86+
"a session that only handshook and listed must have no work session")
87+
}
88+
89+
// No client session in the context (internal calls, tests, odd transports) must
90+
// be a no-op rather than a panic.
91+
func TestWorkSession_NoClientSessionIsNoop(t *testing.T) {
92+
p, _ := newWorkSessionProxy(t)
93+
assert.Empty(t, p.markWorkIfToolCall(context.Background(), mcp.MethodToolsCall))
94+
}

internal/server/workspace.go

Lines changed: 43 additions & 2 deletions
Original file line numberDiff line numberDiff line change
@@ -167,14 +167,55 @@ func (p *MCPProxyServer) markSessionWorked(ctx context.Context, sessionID string
167167
principal := principalFromContext(ctx)
168168

169169
return p.sessionStore.EnsurePersisted(sessionID, func(info *SessionInfo) string {
170-
if p.mainServer == nil || p.mainServer.runtime == nil {
170+
resolve := p.resolveWorkSession()
171+
if resolve == nil {
171172
return ""
172173
}
173-
return p.mainServer.runtime.ResolveWorkSession(runtime.WorkSessionIdentity{
174+
return resolve(runtime.WorkSessionIdentity{
174175
Principal: principal,
175176
ClientName: info.ClientName,
176177
ClientVersion: info.ClientVersion,
177178
WorkspaceRoot: info.Workspace,
178179
})
179180
})
180181
}
182+
183+
// resolveWorkSession is the runtime's work-session tracker, or the stub a test
184+
// injected in its place. Nil when neither is available.
185+
func (p *MCPProxyServer) resolveWorkSession() func(runtime.WorkSessionIdentity) string {
186+
if p.workSessionResolver != nil {
187+
return p.workSessionResolver
188+
}
189+
if p.mainServer == nil || p.mainServer.runtime == nil {
190+
return nil
191+
}
192+
return p.mainServer.runtime.ResolveWorkSession
193+
}
194+
195+
// markWorkIfToolCall attributes a request to a work session when — and only when
196+
// — it is real work (Spec 082).
197+
//
198+
// This runs from the beforeAny hook, which every MCP method passes through on
199+
// every server instance (direct / code-exec / call-tool all share these hooks).
200+
// That placement is the point: work used to be marked inside individual tool
201+
// handlers, and the built-ins that did not bother — list_registries,
202+
// upstream_servers, quarantine_security — wrote activity with no work session on
203+
// it. The Web UI grouped those orphaned rows under the raw transport session id
204+
// while the rest of the same connection grouped under its work-session id, so a
205+
// single client appeared as two sessions in the picker. Marking here means a new
206+
// built-in cannot reintroduce that by forgetting a call.
207+
//
208+
// A tools/call is the boundary of "work", deliberately: a connection that only
209+
// initializes and lists tools has done nothing, and persisting it would bury the
210+
// user's real sessions under background agents that connect every few minutes,
211+
// do nothing, and leave.
212+
func (p *MCPProxyServer) markWorkIfToolCall(ctx context.Context, method mcp.MCPMethod) string {
213+
if method != mcp.MethodToolsCall {
214+
return ""
215+
}
216+
session := mcpserver.ClientSessionFromContext(ctx)
217+
if session == nil {
218+
return ""
219+
}
220+
return p.markSessionWorked(ctx, session.SessionID())
221+
}

0 commit comments

Comments
 (0)