From bdef38398ed56650208b4d6164f415d0f4b82796 Mon Sep 17 00:00:00 2001 From: Prasanna A P <106952318+Prasanna721@users.noreply.github.com> Date: Tue, 21 Jul 2026 00:30:41 -0700 Subject: [PATCH] share workspace selection across sessions --- apps/mcp/e2e/widgets.test.ts | 14 ++++++++- apps/mcp/src/server/agent.ts | 29 ++++++++++++++----- apps/mcp/src/server/auth/index.test.ts | 6 +++- apps/mcp/src/server/auth/index.ts | 5 ++++ apps/mcp/src/server/index.ts | 5 +++- apps/mcp/src/server/tools/guided-save.ts | 2 +- apps/mcp/src/server/tools/select-workspace.ts | 2 +- apps/mcp/src/server/tools/set-active-tag.ts | 2 +- apps/mcp/src/server/tools/types.ts | 6 ++-- apps/mcp/src/server/tools/upload-file.ts | 2 +- apps/mcp/src/server/tools/who-am-i.ts | 2 +- apps/mcp/src/server/workspace-state.ts | 13 +++++++++ apps/mcp/src/shared/types.ts | 1 + apps/mcp/src/widget/hooks/useApp.ts | 3 +- apps/mcp/src/widget/views/Picker.tsx | 15 +++++++++- apps/mcp/src/widget/views/Save.tsx | 15 +++++++++- apps/mcp/src/widget/views/Upload.tsx | 16 +++++++++- apps/mcp/wrangler.jsonc | 8 +++++ 18 files changed, 123 insertions(+), 23 deletions(-) create mode 100644 apps/mcp/src/server/workspace-state.ts diff --git a/apps/mcp/e2e/widgets.test.ts b/apps/mcp/e2e/widgets.test.ts index 0966a449..52951b29 100644 --- a/apps/mcp/e2e/widgets.test.ts +++ b/apps/mcp/e2e/widgets.test.ts @@ -4,6 +4,7 @@ import { callTool, connect, type Session, + textOf, } from "./helpers" describe.skipIf(!OAUTH_CREDENTIALS_AVAILABLE)( @@ -42,7 +43,7 @@ describe.skipIf(!OAUTH_CREDENTIALS_AVAILABLE)( ).toBe(true) }) - it("sets an active workspace only from the visible list", async () => { + it("shares the selected workspace across MCP transport sessions", async () => { const picker = await callTool(session.client, "select-workspace") const pickerContent = picker.structuredContent as { containerTags?: Array<{ containerTag: string }> @@ -58,6 +59,17 @@ describe.skipIf(!OAUTH_CREDENTIALS_AVAILABLE)( view: "confirmation", containerTag: firstTag, }) + + const separateSession = await connect() + try { + const identity = await callTool(separateSession.client, "whoAmI") + expect(identity.isError).toBeFalsy() + expect(JSON.parse(textOf(identity))).toMatchObject({ + activeWorkspace: firstTag, + }) + } finally { + await separateSession.close() + } }) it("loads guided-save writable choices on demand", async () => { diff --git a/apps/mcp/src/server/agent.ts b/apps/mcp/src/server/agent.ts index 53f0987c..69c281c7 100644 --- a/apps/mcp/src/server/agent.ts +++ b/apps/mcp/src/server/agent.ts @@ -9,9 +9,11 @@ import { registerProfileResource } from "./resources/profile" import { registerWidgetResource } from "./resources/widget" import { registerAllTools } from "./tools" import { errorResult } from "./tools/types" +import type { WorkspaceState } from "./workspace-state" type Env = { MCP_SERVER: DurableObjectNamespace + WORKSPACE_STATE: DurableObjectNamespace API_URL?: string } @@ -49,12 +51,11 @@ export class SupermemoryMCP extends McpAgent { getSession: () => this.getSession(), resolveContainerTag: (explicit?: string) => this.resolveContainerTag(explicit), - storage: { - get: (key: string) => this.ctx.storage.get(key), - put: (key: string, value: T) => this.ctx.storage.put(key, value), - }, + getActiveContainerTag: () => this.getActiveContainerTag(), + setActiveContainerTag: (containerTag: string) => + this.setActiveContainerTag(containerTag), getClientInfo: () => this.clientInfo, - getMcpSessionId: () => this.ctx.id.name ?? "unknown", + getMcpSessionId: () => this.getSessionId(), errorResult, } @@ -83,10 +84,24 @@ export class SupermemoryMCP extends McpAgent { private async resolveContainerTag( explicit?: string, ): Promise { + if (this.props?.containerTag) return this.props.containerTag if (explicit) return explicit - const activeTag = await this.ctx.storage.get("activeContainerTag") + const activeTag = await this.getActiveContainerTag() if (activeTag) return activeTag - return this.props?.containerTag + return undefined + } + + private workspaceState() { + const key = `${this.props.userId}:${this.props.organizationId ?? "default"}` + return this.env.WORKSPACE_STATE.getByName(key) + } + + private getActiveContainerTag(): Promise { + return this.workspaceState().getActiveContainerTag() + } + + private setActiveContainerTag(containerTag: string): Promise { + return this.workspaceState().setActiveContainerTag(containerTag) } private getSession() { diff --git a/apps/mcp/src/server/auth/index.test.ts b/apps/mcp/src/server/auth/index.test.ts index 4ec23035..cf3c98f2 100644 --- a/apps/mcp/src/server/auth/index.test.ts +++ b/apps/mcp/src/server/auth/index.test.ts @@ -47,7 +47,11 @@ describe("MCP authentication", () => { await expect( validateOAuthToken(token, API_URL, MCP_RESOURCE, keySet), - ).resolves.toEqual({ userId: "user_test", bearerToken: token }) + ).resolves.toEqual({ + userId: "user_test", + organizationId: "org_test", + bearerToken: token, + }) expect(fetchSpy).not.toHaveBeenCalled() }) diff --git a/apps/mcp/src/server/auth/index.ts b/apps/mcp/src/server/auth/index.ts index 58a1e3b0..393b8e2a 100644 --- a/apps/mcp/src/server/auth/index.ts +++ b/apps/mcp/src/server/auth/index.ts @@ -5,6 +5,7 @@ const FETCH_TIMEOUT_MS = 30_000 export interface AuthUser { userId: string + organizationId?: string bearerToken: string } @@ -66,6 +67,10 @@ export async function validateOAuthToken( } return { userId: payload.sub, + organizationId: + typeof payload.organization_id === "string" + ? payload.organization_id + : undefined, bearerToken: token, } } catch (error) { diff --git a/apps/mcp/src/server/index.ts b/apps/mcp/src/server/index.ts index 746d78c5..84f56552 100644 --- a/apps/mcp/src/server/index.ts +++ b/apps/mcp/src/server/index.ts @@ -4,9 +4,11 @@ import type { ContentfulStatusCode } from "hono/utils/http-status" import type { Props } from "../shared/types" import { SupermemoryMCP } from "./agent" import { validateOAuthToken } from "./auth" +import { WorkspaceState } from "./workspace-state" type Bindings = { MCP_SERVER: DurableObjectNamespace + WORKSPACE_STATE: DurableObjectNamespace API_URL?: string MCP_RESOURCE?: string } @@ -154,6 +156,7 @@ async function handleMcpRequest( ...c.executionCtx, props: { userId: authUser.userId, + organizationId: authUser.organizationId, bearerToken: authUser.bearerToken, containerTag, } satisfies Props, @@ -174,6 +177,6 @@ app.all("/mcp/*", async (c) => { return handleMcpRequest(c) }) -export { SupermemoryMCP } +export { SupermemoryMCP, WorkspaceState } export default app diff --git a/apps/mcp/src/server/tools/guided-save.ts b/apps/mcp/src/server/tools/guided-save.ts index 1cceceff..4c52f6ce 100644 --- a/apps/mcp/src/server/tools/guided-save.ts +++ b/apps/mcp/src/server/tools/guided-save.ts @@ -20,7 +20,7 @@ export function register(deps: ToolDeps) { try { const prefill = (args as { prefill?: string }).prefill const [activeTag, tags, session] = await Promise.all([ - deps.storage.get("activeContainerTag"), + deps.getActiveContainerTag(), deps.getClient().listContainerTags(), deps.getSession(), ]) diff --git a/apps/mcp/src/server/tools/select-workspace.ts b/apps/mcp/src/server/tools/select-workspace.ts index 191f4d50..51360982 100644 --- a/apps/mcp/src/server/tools/select-workspace.ts +++ b/apps/mcp/src/server/tools/select-workspace.ts @@ -20,7 +20,7 @@ export function register(deps: ToolDeps) { const [tags, session, activeTag] = await Promise.all([ client.listContainerTags(), deps.getSession(), - deps.storage.get("activeContainerTag"), + deps.getActiveContainerTag(), ]) const assignedTags = effectiveContainerTagAccess( tags.map((tag) => tag.containerTag), diff --git a/apps/mcp/src/server/tools/set-active-tag.ts b/apps/mcp/src/server/tools/set-active-tag.ts index a8e653f8..397589e2 100644 --- a/apps/mcp/src/server/tools/set-active-tag.ts +++ b/apps/mcp/src/server/tools/set-active-tag.ts @@ -28,7 +28,7 @@ export function register(deps: ToolDeps) { new Error(`No access to container tag '${containerTag}'.`), ) } - await deps.storage.put("activeContainerTag", containerTag) + await deps.setActiveContainerTag(containerTag) const sc: ViewMessage = { view: "confirmation", containerTag, diff --git a/apps/mcp/src/server/tools/types.ts b/apps/mcp/src/server/tools/types.ts index f14f1842..18b6e5d4 100644 --- a/apps/mcp/src/server/tools/types.ts +++ b/apps/mcp/src/server/tools/types.ts @@ -10,10 +10,8 @@ export interface ToolDeps { getClient: (containerTag?: string) => SupermemoryClient getSession: () => Promise resolveContainerTag: (explicit?: string) => Promise - storage: { - get: (key: string) => Promise - put: (key: string, value: T) => Promise - } + getActiveContainerTag: () => Promise + setActiveContainerTag: (containerTag: string) => Promise getClientInfo: () => { name: string; version?: string } | null getMcpSessionId: () => string errorResult: (error: unknown) => { diff --git a/apps/mcp/src/server/tools/upload-file.ts b/apps/mcp/src/server/tools/upload-file.ts index 0c0d1436..30be4c23 100644 --- a/apps/mcp/src/server/tools/upload-file.ts +++ b/apps/mcp/src/server/tools/upload-file.ts @@ -16,7 +16,7 @@ export function register(deps: ToolDeps) { async () => { try { const [activeTag, tags, session] = await Promise.all([ - deps.storage.get("activeContainerTag"), + deps.getActiveContainerTag(), deps.getClient().listContainerTags(), deps.getSession(), ]) diff --git a/apps/mcp/src/server/tools/who-am-i.ts b/apps/mcp/src/server/tools/who-am-i.ts index 0a6f0eb8..8281f5ae 100644 --- a/apps/mcp/src/server/tools/who-am-i.ts +++ b/apps/mcp/src/server/tools/who-am-i.ts @@ -11,7 +11,7 @@ export function register(deps: ToolDeps) { try { const [session, activeTag] = await Promise.all([ deps.getSession(), - deps.storage.get("activeContainerTag"), + deps.getActiveContainerTag(), ]) return { content: [ diff --git a/apps/mcp/src/server/workspace-state.ts b/apps/mcp/src/server/workspace-state.ts new file mode 100644 index 00000000..8ad756a2 --- /dev/null +++ b/apps/mcp/src/server/workspace-state.ts @@ -0,0 +1,13 @@ +import { DurableObject } from "cloudflare:workers" + +const ACTIVE_CONTAINER_TAG_KEY = "activeContainerTag" + +export class WorkspaceState extends DurableObject { + async getActiveContainerTag(): Promise { + return this.ctx.storage.get(ACTIVE_CONTAINER_TAG_KEY) + } + + async setActiveContainerTag(containerTag: string): Promise { + await this.ctx.storage.put(ACTIVE_CONTAINER_TAG_KEY, containerTag) + } +} diff --git a/apps/mcp/src/shared/types.ts b/apps/mcp/src/shared/types.ts index 3f67302d..c46eb2d8 100644 --- a/apps/mcp/src/shared/types.ts +++ b/apps/mcp/src/shared/types.ts @@ -119,6 +119,7 @@ export type ViewName = ViewMessage["view"] // Auth context passed from the OAuth/API-key middleware into the McpAgent via ctx.props. export type Props = { userId: string + organizationId?: string bearerToken: string containerTag?: string } diff --git a/apps/mcp/src/widget/hooks/useApp.ts b/apps/mcp/src/widget/hooks/useApp.ts index 8f1af503..2c5ff25a 100644 --- a/apps/mcp/src/widget/hooks/useApp.ts +++ b/apps/mcp/src/widget/hooks/useApp.ts @@ -56,9 +56,10 @@ export function useApp() { }, /** Update ambient model context. Model sees on next user message. */ - updateContext(text: string) { + updateContext(text: string, structuredContent?: Record) { return app.updateModelContext({ content: [{ type: "text", text }], + structuredContent, }) }, diff --git a/apps/mcp/src/widget/views/Picker.tsx b/apps/mcp/src/widget/views/Picker.tsx index 45c69605..48eae948 100644 --- a/apps/mcp/src/widget/views/Picker.tsx +++ b/apps/mcp/src/widget/views/Picker.tsx @@ -28,7 +28,7 @@ export function Picker({ onAdvance, onError, }: Props) { - const { callTool } = useApp() + const { callTool, updateContext } = useApp() const log = useLog() const [pending, setPending] = useState(null) const [query, setQuery] = useState("") @@ -55,6 +55,19 @@ export function Picker({ onError(result.error ?? "Failed to set active workspace") return } + try { + await updateContext( + `The user selected the Supermemory workspace "${containerTag}". Use this as the active workspace for future Supermemory actions unless the user selects another workspace.`, + { + supermemory: { + activeWorkspace: containerTag, + lastAction: "workspace-selected", + }, + }, + ) + } catch (error) { + log("warning", `[picker] model context update failed: ${error}`) + } onAdvance(result.data) } diff --git a/apps/mcp/src/widget/views/Save.tsx b/apps/mcp/src/widget/views/Save.tsx index 0dbd9b16..2ba5c7e6 100644 --- a/apps/mcp/src/widget/views/Save.tsx +++ b/apps/mcp/src/widget/views/Save.tsx @@ -27,7 +27,7 @@ export function Save({ onAdvance, onError, }: Props) { - const { callTool } = useApp() + const { callTool, updateContext } = useApp() const log = useLog() const [content, setContent] = useState(prefill ?? "") const [selectedTag, setSelectedTag] = useState( @@ -63,6 +63,19 @@ export function Save({ onError(result.error ?? "Failed to save memory") return } + try { + await updateContext( + `The user saved a memory to the Supermemory workspace "${selectedTag}" from the interactive widget.`, + { + supermemory: { + activeWorkspace: selectedTag, + lastAction: "memory-saved", + }, + }, + ) + } catch (error) { + log("warning", `[save] model context update failed: ${error}`) + } onAdvance(result.data) } diff --git a/apps/mcp/src/widget/views/Upload.tsx b/apps/mcp/src/widget/views/Upload.tsx index dda2e10f..6b0af03c 100644 --- a/apps/mcp/src/widget/views/Upload.tsx +++ b/apps/mcp/src/widget/views/Upload.tsx @@ -30,7 +30,7 @@ function formatFileSize(bytes: number): string { const ACCEPT = ".txt,.pdf,.png,.jpg,.jpeg,.mp4" export function Upload({ activeTag, writableTags, onAdvance, onError }: Props) { - const { callTool } = useApp() + const { callTool, updateContext } = useApp() const log = useLog() const [file, setFile] = useState(null) const [selectedTag, setSelectedTag] = useState( @@ -62,6 +62,20 @@ export function Upload({ activeTag, writableTags, onAdvance, onError }: Props) { onError(result.error ?? "Upload failed") return } + try { + await updateContext( + `The user uploaded "${file.name}" to the Supermemory workspace "${selectedTag}" from the interactive widget.`, + { + supermemory: { + activeWorkspace: selectedTag, + lastAction: "file-uploaded", + fileName: file.name, + }, + }, + ) + } catch (error) { + log("warning", `[upload] model context update failed: ${error}`) + } onAdvance(result.data) } catch (err) { log("error", `[upload] threw: ${err}`) diff --git a/apps/mcp/wrangler.jsonc b/apps/mcp/wrangler.jsonc index a7429795..8a85739f 100644 --- a/apps/mcp/wrangler.jsonc +++ b/apps/mcp/wrangler.jsonc @@ -27,6 +27,10 @@ { "name": "MCP_SERVER", "class_name": "SupermemoryMCP" + }, + { + "name": "WORKSPACE_STATE", + "class_name": "WorkspaceState" } ] }, @@ -35,6 +39,10 @@ { "tag": "v1", "new_sqlite_classes": ["SupermemoryMCP"] + }, + { + "tag": "v2", + "new_sqlite_classes": ["WorkspaceState"] } ],