Skip to content

Commit 2d9399a

Browse files
Sg312Sg312
committed
fix(mothership): provide authorized context for chat titles
1 parent 67776c3 commit 2d9399a

5 files changed

Lines changed: 265 additions & 0 deletions

File tree

‎apps/sim/lib/mothership/chat/title-context.test.ts‎

Lines changed: 123 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -0,0 +1,123 @@
1+
import { beforeEach, describe, expect, it, vi } from 'vitest'
2+
import { buildChatTitleContext } from '@/lib/mothership/chat/title-context'
3+
4+
const { getWorkspace, listWorkspaces } = vi.hoisted(() => ({
5+
getWorkspace: vi.fn(),
6+
listWorkspaces: vi.fn(),
7+
}))
8+
vi.mock('@/lib/workspaces/application/get-public-workspace', () => ({
9+
getPublicWorkspace: { execute: getWorkspace },
10+
}))
11+
vi.mock('@/lib/workspaces/application/list-organization-workspaces', () => ({
12+
listOrganizationWorkspaces: { execute: listWorkspaces },
13+
}))
14+
const inventory = {
15+
workflows: [{ id: 'workflow', name: 'Invoice approval', deployed: true }],
16+
tables: [{ id: 'table', name: 'Customers' }],
17+
knowledgeBases: [],
18+
files: [{ path: 'invoice.csv' }],
19+
skills: [],
20+
customTools: [],
21+
mcpServers: [],
22+
credentials: [{ id: 'credential', name: 'SECRET CREDENTIAL' }],
23+
secrets: ['SECRET NAME'],
24+
truncated: [],
25+
}
26+
27+
describe('buildChatTitleContext', () => {
28+
beforeEach(() => {
29+
vi.clearAllMocks()
30+
getWorkspace.mockResolvedValue({ workspace: { name: 'Finance' } })
31+
})
32+
33+
it('reuses prepared or recovered inventory names after current workspace authorization', async () => {
34+
const result = await buildChatTitleContext({
35+
userId: 'user',
36+
workspaceId: 'workspace',
37+
inventory,
38+
})
39+
expect(JSON.parse(result!)).toEqual({
40+
scope: 'workspace',
41+
workspaceName: 'Finance',
42+
partial: true,
43+
resources: { workflows: ['Invoice approval'], tables: ['Customers'], files: ['invoice.csv'] },
44+
})
45+
expect(result).not.toContain('SECRET')
46+
expect(getWorkspace).toHaveBeenCalledWith({
47+
principal: expect.objectContaining({
48+
subjectUserId: 'user',
49+
workspaceId: 'workspace',
50+
audience: 'sim:workspaces',
51+
}),
52+
input: { workspaceId: 'workspace' },
53+
})
54+
expect(listWorkspaces).not.toHaveBeenCalled()
55+
})
56+
57+
it('uses a bounded authorized org listing and never fans out into workspace inventories', async () => {
58+
listWorkspaces.mockResolvedValue({ workspaces: [{ name: 'Marketing', id: 'hidden-id' }] })
59+
const result = await buildChatTitleContext({
60+
userId: 'user',
61+
organizationId: 'org',
62+
chatId: 'chat',
63+
inventory,
64+
})
65+
expect(JSON.parse(result!)).toEqual({
66+
scope: 'organization',
67+
partial: true,
68+
resources: { workspaces: ['Marketing'] },
69+
})
70+
expect(listWorkspaces).toHaveBeenCalledWith({
71+
principal: expect.objectContaining({
72+
subjectUserId: 'user',
73+
organizationId: 'org',
74+
resourceScope: { chatId: 'chat' },
75+
audience: 'sim:workspaces',
76+
}),
77+
input: { organizationId: 'org', limit: 12 },
78+
})
79+
expect(getWorkspace).not.toHaveBeenCalled()
80+
})
81+
82+
it('omits metadata on revoked access, including previously prepared inventory', async () => {
83+
getWorkspace.mockRejectedValueOnce(new Error('forbidden'))
84+
expect(
85+
await buildChatTitleContext({ userId: 'user', workspaceId: 'workspace', inventory })
86+
).toBeUndefined()
87+
})
88+
89+
it('bounds escaped names and excludes all unprojected fields', async () => {
90+
const names = Array.from({ length: 100 }, () => ({
91+
id: 'id',
92+
name: '\u0000'.repeat(500),
93+
deployed: true,
94+
}))
95+
const result = await buildChatTitleContext({
96+
userId: 'user',
97+
workspaceId: 'workspace',
98+
inventory: {
99+
...inventory,
100+
workflows: names,
101+
tables: names,
102+
knowledgeBases: names,
103+
body: 'PRIVATE BODY',
104+
inputSchema: { secret: 'PRIVATE SCHEMA' },
105+
},
106+
})
107+
expect(result!.length).toBeLessThanOrEqual(6000)
108+
expect(() => JSON.parse(result!)).not.toThrow()
109+
expect(result).not.toContain('PRIVATE')
110+
})
111+
112+
it('does not invent authority for missing owners or ambiguous scope', async () => {
113+
expect(await buildChatTitleContext({ workspaceId: 'workspace', inventory })).toBeUndefined()
114+
expect(
115+
await buildChatTitleContext({
116+
userId: 'user',
117+
workspaceId: 'workspace',
118+
organizationId: 'org',
119+
})
120+
).toBeUndefined()
121+
expect(getWorkspace).not.toHaveBeenCalled()
122+
})
123+
})

‎apps/sim/lib/mothership/chat/title-context.ts‎

Lines changed: 93 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -0,0 +1,93 @@
1+
import { createLogger } from '@sim/logger'
2+
import { truncate } from '@sim/utils/string'
3+
import {
4+
createCopilotChatPrincipal,
5+
createTrustedOrganizationCopilotPrincipal,
6+
} from '@/lib/mothership/auth/application-delegation'
7+
import { ChatPayloadSchema } from '@/lib/mothership/generated/protocol'
8+
import { getPublicWorkspace } from '@/lib/workspaces/application/get-public-workspace'
9+
import { listOrganizationWorkspaces } from '@/lib/workspaces/application/list-organization-workspaces'
10+
11+
const logger = createLogger('ChatTitleContext')
12+
const NAME_LIMIT = 120
13+
const CONTEXT_LIMIT = 6000
14+
15+
/** Optional, names-only orientation for the background title request. */
16+
export async function buildChatTitleContext(input: {
17+
userId?: string
18+
workspaceId?: string
19+
organizationId?: string
20+
chatId?: string
21+
inventory?: unknown
22+
}): Promise<string | undefined> {
23+
if (!input.userId) return undefined
24+
try {
25+
if (input.organizationId && !input.workspaceId && input.chatId) {
26+
const { workspaces } = await listOrganizationWorkspaces.execute({
27+
principal: createTrustedOrganizationCopilotPrincipal(
28+
{
29+
userId: input.userId,
30+
organizationId: input.organizationId,
31+
chatId: input.chatId,
32+
delegationId: `title:${input.chatId}`,
33+
},
34+
{ audience: 'sim:workspaces', ttlMs: 60_000 }
35+
),
36+
input: { organizationId: input.organizationId, limit: 12 },
37+
})
38+
return boundedSnapshot('organization', undefined, {
39+
workspaces: workspaces.map(({ name }) => name),
40+
})
41+
}
42+
if (!input.workspaceId || input.organizationId) return undefined
43+
const { workspace } = await getPublicWorkspace.execute({
44+
principal: createCopilotChatPrincipal(
45+
{ userId: input.userId, workspaceId: input.workspaceId, chatId: input.chatId },
46+
'sim:workspaces'
47+
),
48+
input: { workspaceId: input.workspaceId },
49+
})
50+
const parsed = ChatPayloadSchema.shape.inventory.safeParse(input.inventory)
51+
const inventory = parsed.success ? parsed.data : undefined
52+
return boundedSnapshot('workspace', workspace.name, {
53+
workflows: inventory?.workflows.map(({ name }) => name) ?? [],
54+
tables: inventory?.tables.map(({ name }) => name) ?? [],
55+
knowledgeBases: inventory?.knowledgeBases.map(({ name }) => name) ?? [],
56+
files: inventory?.files.map(({ path }) => path) ?? [],
57+
skills: inventory?.skills.map(({ name }) => name) ?? [],
58+
customTools: inventory?.customTools.map(({ title }) => title) ?? [],
59+
mcpServers: inventory?.mcpServers.map(({ name }) => name) ?? [],
60+
})
61+
} catch {
62+
logger.warn('Title context metadata unavailable')
63+
return undefined
64+
}
65+
}
66+
67+
/** Bound the serialized size too, since names may contain JSON escape characters. */
68+
function boundedSnapshot(
69+
scope: string,
70+
workspaceName: string | undefined,
71+
groups: Record<string, string[]>
72+
): string {
73+
const snapshot = {
74+
scope,
75+
...(workspaceName ? { workspaceName: truncate(workspaceName, NAME_LIMIT) } : {}),
76+
partial: true,
77+
resources: {} as Record<string, string[]>,
78+
}
79+
for (const [kind, names] of Object.entries(groups)) {
80+
if (!names.length) continue
81+
const selected: string[] = []
82+
snapshot.resources[kind] = selected
83+
for (const name of names.slice(0, 12)) {
84+
selected.push(truncate(name, NAME_LIMIT))
85+
if (JSON.stringify(snapshot).length > CONTEXT_LIMIT) {
86+
selected.pop()
87+
break
88+
}
89+
}
90+
if (!selected.length) delete snapshot.resources[kind]
91+
}
92+
return JSON.stringify(snapshot)
93+
}

‎apps/sim/lib/mothership/generated/protocol.ts‎

Lines changed: 2 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -367,6 +367,8 @@ export type ForkChatResponse = z.infer<typeof ForkChatResponse>;
367367
/** POST /api/generate-chat-title */
368368
export const TitleRequest = z.strictObject({
369369
message: z.string().min(1),
370+
/** Bounded, authorized scope and resource names for disambiguating the user's topic. */
371+
context: z.string().max(6_000).optional(),
370372
/** Enterprise BYOK pins the title call because it reads user content. */
371373
byokApiKey: z.string().optional(),
372374
/** Optional metering identity for callers without a persisted chat yet. */

‎apps/sim/lib/mothership/request/lifecycle/start.test.ts‎

Lines changed: 40 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -32,6 +32,7 @@ const {
3232
releasePendingChatStream,
3333
unregisterActiveStream,
3434
fetchGo,
35+
buildChatTitleContext,
3536
} = vi.hoisted(() => ({
3637
runCopilotLifecycle: vi.fn(),
3738
createRunSegment: vi.fn(),
@@ -49,6 +50,7 @@ const {
4950
releasePendingChatStream: vi.fn(),
5051
unregisterActiveStream: vi.fn(),
5152
fetchGo: vi.fn(),
53+
buildChatTitleContext: vi.fn().mockResolvedValue(undefined),
5254
}))
5355

5456
const BILLING_ATTRIBUTION = {
@@ -75,6 +77,8 @@ vi.mock('@/lib/mothership/request/session/controller-lease', async (original) =>
7577
assertChatStreamLease: vi.fn().mockResolvedValue(undefined),
7678
}))
7779

80+
vi.mock('@/lib/mothership/chat/title-context', () => ({ buildChatTitleContext }))
81+
7882
vi.mock('@/lib/mothership/request/lifecycle/run', () => ({
7983
runCopilotLifecycle,
8084
}))
@@ -400,6 +404,42 @@ describe('createSSEStream terminal error handling', () => {
400404
}
401405
)
402406

407+
it('starts the agent while title metadata is pending and forwards recovered inventory', async () => {
408+
let resolveContext!: (value: string) => void
409+
buildChatTitleContext.mockImplementationOnce(
410+
() =>
411+
new Promise<string>((resolve) => {
412+
resolveContext = resolve
413+
})
414+
)
415+
fetchGo.mockResolvedValueOnce(Response.json({ title: 'Finance planning' }))
416+
const recoveredInventory = { workflows: [{ name: 'Invoice approval' }] }
417+
const stream = createSSEStream({
418+
requestPayload: { message: 'Plan', inventory: recoveredInventory },
419+
userId: 'user-1',
420+
workspaceId: '22222222-2222-4222-8222-222222222222',
421+
chatId: '11111111-1111-4111-8111-111111111111',
422+
streamId: 'stream-title',
423+
executionId: 'exec-title',
424+
runId: 'run-title',
425+
currentChat: null,
426+
message: 'Plan',
427+
titleModel: 'gpt-5.4',
428+
requestId: 'req-title',
429+
orchestrateOptions: { userId: 'user-1' },
430+
})
431+
await vi.waitFor(() => expect(buildChatTitleContext).toHaveBeenCalled())
432+
expect(runCopilotLifecycle).toHaveBeenCalled()
433+
expect(fetchGo).not.toHaveBeenCalled()
434+
expect(buildChatTitleContext).toHaveBeenCalledWith(
435+
expect.objectContaining({ inventory: recoveredInventory })
436+
)
437+
resolveContext('{"workspaceName":"Finance"}')
438+
await vi.waitFor(() => expect(fetchGo).toHaveBeenCalled())
439+
expect(JSON.parse(fetchGo.mock.calls[0][1].body).context).toBe('{"workspaceName":"Finance"}')
440+
await drainStream(stream)
441+
})
442+
403443
it('finishes a pre-admission Stop as cancelled without calling the agent or title model', async () => {
404444
createRunSegment.mockResolvedValueOnce({ status: 'cancelled' })
405445
await drainStream(

‎apps/sim/lib/mothership/request/lifecycle/start.ts‎

Lines changed: 7 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -13,6 +13,7 @@ import {
1313
} from '@/lib/billing/core/billing-attribution'
1414
import { isHosted } from '@/lib/core/config/env-flags'
1515
import { createRunSegment, recordRunBillingAdmission } from '@/lib/mothership/async-runs/repository'
16+
import { buildChatTitleContext } from '@/lib/mothership/chat/title-context'
1617
import { publishChatStatusChanged } from '@/lib/mothership/chat-status'
1718
import {
1819
MothershipStreamV1EventType,
@@ -304,6 +305,7 @@ export function createSSEStream(params: StreamingOrchestrationParams): ReadableS
304305
}
305306

306307
fireTitleGeneration({
308+
inventory: requestPayload.inventory,
307309
chatId,
308310
currentChat,
309311
userId,
@@ -520,6 +522,7 @@ export function createSSEStream(params: StreamingOrchestrationParams): ReadableS
520522
// Title generation (fire-and-forget side effect)
521523

522524
function fireTitleGeneration(params: {
525+
inventory?: unknown
523526
chatId?: string
524527
currentChat: CurrentChatSummary
525528
userId?: string
@@ -551,6 +554,7 @@ function fireTitleGeneration(params: {
551554
if (!chatId || currentChat?.title) return
552555

553556
requestChatTitle({
557+
inventory: params.inventory,
554558
chatId,
555559
message,
556560
model: titleModel,
@@ -587,6 +591,7 @@ function fireTitleGeneration(params: {
587591

588592
/** Requests a title through the shared Assistant backend and its attributed billing protocol. */
589593
export async function requestChatTitle(params: {
594+
inventory?: unknown
590595
chatId?: string
591596
message: string
592597
model: string
@@ -642,13 +647,15 @@ export async function requestChatTitle(params: {
642647
const mothershipBaseURL = await getMothershipBaseURL({ userId })
643648
// Title reads the user's message content, so an enterprise chat pins its key here too.
644649
const byokApiKey = await resolveEnterpriseByokKey(workspaceId)
650+
const context = await buildChatTitleContext(params)
645651
const response = await fetchGo(`${mothershipBaseURL}/api/generate-chat-title`, {
646652
method: 'POST',
647653
signal,
648654
headers,
649655
body: JSON.stringify(
650656
TitleRequest.parse({
651657
message,
658+
...(context ? { context } : {}),
652659
...(workspaceId ? { workspaceId } : {}),
653660
...(organizationId ? { organizationId, chatId } : {}),
654661
...(userId ? { userId } : {}),

0 commit comments

Comments
 (0)