diff --git a/sdk/packages/core/src/hub/server/boundary.test.ts b/sdk/packages/core/src/hub/server/boundary.test.ts index 9bf11dddce..fcf95c3ff4 100644 --- a/sdk/packages/core/src/hub/server/boundary.test.ts +++ b/sdk/packages/core/src/hub/server/boundary.test.ts @@ -13,6 +13,7 @@ import { requestToolApproval, } from "./handlers/approval-handlers"; import { + ensureSessionParticipant, ensureSessionState, type HubTransportContext, } from "./handlers/context"; @@ -859,6 +860,148 @@ describe("HubServerTransport boundaries", () => { expect(readSessionCompactionState).not.toHaveBeenCalled(); }); + it("does not grant compaction sidecar ownership from session attach", async () => { + const state = createSessionCompactionState({ + sourceMessages: [{ role: "user", content: "source" }], + compactedMessages: [{ role: "user", content: "summary" }], + conversationId: "session-1", + }); + const readSessionCompactionState = vi.fn().mockResolvedValue(state); + const updateSessionCompactionState = vi + .fn() + .mockResolvedValue({ updated: true }); + const transport = createTransport({ + sessionHost: { + readSessionCompactionState, + updateSessionCompactionState, + }, + }); + const ctx = getContext(transport); + expect(ctx.sessionState.has("session-1")).toBe(false); + + const attachReply = await transport.handleCommand({ + version: "v1", + requestId: "req-attach", + command: "session.attach", + clientId: "viewer-client", + sessionId: "session-1", + }); + + expect(attachReply).toMatchObject({ ok: true }); + expect( + ctx.sessionState.get("session-1")?.createdByClientId, + ).toBeUndefined(); + expect( + ctx.sessionState.get("session-1")?.participants.has("viewer-client"), + ).toBe(true); + + const getReply = await transport.handleCommand({ + version: "v1", + requestId: "req-compact-get", + command: "session.compaction.get", + clientId: "viewer-client", + sessionId: "session-1", + }); + const updateReply = await transport.handleCommand({ + version: "v1", + requestId: "req-compact-update", + command: "session.compaction.update", + clientId: "viewer-client", + sessionId: "session-1", + payload: { state }, + }); + + expect(getReply).toMatchObject({ + ok: false, + error: { code: "session_wrong_client" }, + }); + expect(updateReply).toMatchObject({ + ok: false, + error: { code: "session_wrong_client" }, + }); + expect(readSessionCompactionState).not.toHaveBeenCalled(); + expect(updateSessionCompactionState).not.toHaveBeenCalled(); + }); + + it("clears compaction sidecar ownership when the owner detaches", async () => { + const readSessionCompactionState = vi.fn(); + const transport = createTransport({ + sessionHost: { readSessionCompactionState }, + }); + const ctx = getContext(transport); + ensureSessionState(ctx, "session-1", "owner-client", "creator"); + ensureSessionParticipant(ctx, "session-1", "viewer-client", "participant"); + + const detachReply = await transport.handleCommand({ + version: "v1", + requestId: "req-detach", + command: "session.detach", + clientId: "owner-client", + sessionId: "session-1", + }); + + expect(detachReply).toMatchObject({ ok: true }); + expect( + ctx.sessionState.get("session-1")?.participants.has("viewer-client"), + ).toBe(true); + expect( + ctx.sessionState.get("session-1")?.createdByClientId, + ).toBeUndefined(); + + const getReply = await transport.handleCommand({ + version: "v1", + requestId: "req-compact-get", + command: "session.compaction.get", + clientId: "viewer-client", + sessionId: "session-1", + }); + + expect(getReply).toMatchObject({ + ok: false, + error: { code: "session_wrong_client" }, + }); + expect(readSessionCompactionState).not.toHaveBeenCalled(); + }); + + it("clears compaction sidecar ownership when the owner unregisters", async () => { + const readSessionCompactionState = vi.fn(); + const transport = createTransport({ + sessionHost: { readSessionCompactionState }, + }); + const ctx = getContext(transport); + ensureSessionState(ctx, "session-1", "owner-client", "creator"); + ensureSessionParticipant(ctx, "session-1", "viewer-client", "participant"); + + const unregisterReply = await transport.handleCommand({ + version: "v1", + requestId: "req-unregister", + command: "client.unregister", + clientId: "owner-client", + }); + + expect(unregisterReply).toMatchObject({ ok: true }); + expect( + ctx.sessionState.get("session-1")?.participants.has("viewer-client"), + ).toBe(true); + expect( + ctx.sessionState.get("session-1")?.createdByClientId, + ).toBeUndefined(); + + const getReply = await transport.handleCommand({ + version: "v1", + requestId: "req-compact-get", + command: "session.compaction.get", + clientId: "viewer-client", + sessionId: "session-1", + }); + + expect(getReply).toMatchObject({ + ok: false, + error: { code: "session_wrong_client" }, + }); + expect(readSessionCompactionState).not.toHaveBeenCalled(); + }); + it("returns compaction sidecar state to the server-owned session client", async () => { const state = createSessionCompactionState({ sourceMessages: [{ role: "user", content: "source" }], diff --git a/sdk/packages/core/src/hub/server/handlers/context.ts b/sdk/packages/core/src/hub/server/handlers/context.ts index 5b0f56d617..ddc34bc6d6 100644 --- a/sdk/packages/core/src/hub/server/handlers/context.ts +++ b/sdk/packages/core/src/hub/server/handlers/context.ts @@ -209,3 +209,41 @@ export function ensureSessionState( ctx.sessionState.set(sessionId, state); return state; } + +export function ensureSessionParticipant( + ctx: HubTransportContext, + sessionId: string, + clientId: string, + role: SessionParticipant["role"], + options: { interactive?: boolean } = {}, +): HubSessionState { + const existing = ctx.sessionState.get(sessionId); + if (existing) { + if (options.interactive !== undefined) { + existing.interactive = options.interactive; + } + if (!existing.participants.has(clientId)) { + existing.participants.set(clientId, { + clientId, + attachedAt: Date.now(), + role, + }); + } + return existing; + } + const state: HubSessionState = { + interactive: options.interactive ?? true, + participants: new Map([ + [ + clientId, + { + clientId, + attachedAt: Date.now(), + role, + }, + ], + ]), + }; + ctx.sessionState.set(sessionId, state); + return state; +} diff --git a/sdk/packages/core/src/hub/server/handlers/session-handlers.ts b/sdk/packages/core/src/hub/server/handlers/session-handlers.ts index 70644350e1..57aa93b710 100644 --- a/sdk/packages/core/src/hub/server/handlers/session-handlers.ts +++ b/sdk/packages/core/src/hub/server/handlers/session-handlers.ts @@ -20,6 +20,7 @@ import { toHubSessionRecord } from "../hub-session-records"; import { cancelPendingCapabilityRequests } from "./capability-handlers"; import { asPlainRecord, + ensureSessionParticipant, ensureSessionState, errorReply, extractSessionId, @@ -630,23 +631,29 @@ export async function handleSessionAttach( "session.attach requires a session id", ); } - ensureSessionState( + const session = await readHubSessionRecord(ctx, sessionId); + if (!session) { + return errorReply( + envelope, + "session_not_found", + `Unknown session: ${sessionId}`, + ); + } + ensureSessionParticipant( ctx, sessionId, envelope.clientId?.trim() || "hub-client", "participant", ); - const session = await readHubSessionRecord(ctx, sessionId); - if (session) { - ctx.publish(ctx.buildEvent("session.attached", { session }, sessionId)); - } - return session - ? okReply(envelope, { session }) - : errorReply( - envelope, - "session_not_found", - `Unknown session: ${sessionId}`, - ); + const attachedSession = await readHubSessionRecord(ctx, sessionId); + ctx.publish( + ctx.buildEvent( + "session.attached", + { session: attachedSession ?? session }, + sessionId, + ), + ); + return okReply(envelope, { session: attachedSession ?? session }); } export async function handleSessionDetach( @@ -662,12 +669,11 @@ export async function handleSessionDetach( ); } const clientId = envelope.clientId?.trim() || "hub-client"; - const ownerClientId = getCapabilityOwnerClientId(ctx, sessionId) ?? clientId; const state = ctx.sessionState.get(sessionId); if (state) { state.participants.delete(clientId); if (state.createdByClientId === clientId) { - state.createdByClientId = ownerClientId; + state.createdByClientId = undefined; } if (state.participants.size === 0) { ctx.sessionState.delete(sessionId); diff --git a/sdk/packages/core/src/hub/server/hub-server-transport.ts b/sdk/packages/core/src/hub/server/hub-server-transport.ts index 51695d6a5a..161122eadb 100644 --- a/sdk/packages/core/src/hub/server/hub-server-transport.ts +++ b/sdk/packages/core/src/hub/server/hub-server-transport.ts @@ -553,6 +553,9 @@ export class HubServerTransport implements NativeHubTransport { private detachClientFromSessions(clientId: string): void { for (const [sessionId, state] of this.sessionState.entries()) { state.participants.delete(clientId); + if (state.createdByClientId === clientId) { + state.createdByClientId = undefined; + } if (state.participants.size === 0) { this.sessionState.delete(sessionId); } diff --git a/sdk/packages/core/src/hub/server/hub-session-records.ts b/sdk/packages/core/src/hub/server/hub-session-records.ts index 6372143419..3e2f8bb96f 100644 --- a/sdk/packages/core/src/hub/server/hub-session-records.ts +++ b/sdk/packages/core/src/hub/server/hub-session-records.ts @@ -7,7 +7,7 @@ import type { SessionAccumulatedUsage } from "../../runtime/host/runtime-host"; import type { SessionRecord as LocalSessionRecord } from "../../types/sessions"; export type HubSessionState = { - createdByClientId: string; + createdByClientId?: string; interactive: boolean; participants: Map; };