Preserve chat recovery and unify cancellation by mzxchandra · Pull Request #9142 · simstudioai/sim · GitHub
Skip to content
Draft
Show file tree
Hide file tree
Changes from all commits
Commits
File filter

Filter by extension

Filter by extension

Conversations
Failed to load comments.
Loading
Jump to
Jump to file
Failed to load files.
Loading
Diff view
Diff view
4 changes: 3 additions & 1 deletion apps/sim/app/api/copilot/api-keys/validate/route.ts
2 changes: 1 addition & 1 deletion apps/sim/app/api/mothership/integrations/catalog/route.ts
Original file line number Diff line number Diff line change
Expand Up @@ -13,7 +13,7 @@ import {

export const POST = defineInternalJsonRoute({
contract: readIntegrationCatalogContract,
auth: internalCopilotAuth(INTEGRATION_CATALOG_AUDIENCE, { organization: true }),
auth: internalCopilotAuth(INTEGRATION_CATALOG_AUDIENCE, { organization: true, personal: true }),
operation: readIntegrationCatalogOperation,
rateLimit: internalRateLimits.none({
reason: 'Private bounded catalog discovery rechecks current authorization on every request.',
Expand Down
2 changes: 1 addition & 1 deletion apps/sim/app/api/mothership/memories/access/route.ts
Original file line number Diff line number Diff line change
Expand Up @@ -13,7 +13,7 @@ import {

export const POST = defineInternalJsonRoute({
contract: resolveMemoryNotesAccessContract,
auth: internalCopilotAuth(MEMORY_NOTES_ACCESS_AUDIENCE, { organization: true }),
auth: internalCopilotAuth(MEMORY_NOTES_ACCESS_AUDIENCE, { organization: true, personal: true }),
operation: resolveMemoryNotesAccessOperation,
rateLimit: internalRateLimits.none({
reason:
Expand Down
2 changes: 1 addition & 1 deletion apps/sim/app/api/mothership/memory/scope/route.ts
Original file line number Diff line number Diff line change
Expand Up @@ -13,7 +13,7 @@ import {

export const POST = defineInternalJsonRoute({
contract: readMemoryScopeContract,
auth: internalCopilotAuth(MEMORY_SCOPE_AUDIENCE, { organization: true }),
auth: internalCopilotAuth(MEMORY_SCOPE_AUDIENCE, { organization: true, personal: true }),
operation: readMemoryScopeOperation,
rateLimit: internalRateLimits.none({
reason: 'Authenticated worker resolves its current private chat memory scope.',
Expand Down
36 changes: 25 additions & 11 deletions apps/sim/app/api/mothership/runs/control/route.test.ts
Original file line number Diff line number Diff line change
Expand Up @@ -29,7 +29,7 @@ import { POST } from '@/app/api/mothership/runs/control/route'
const mocks = {
run: mothershipAsyncRunsMockFns.mockGetLatestRunForStream,
stopped: mothershipAsyncRunsMockFns.mockIsRunStopRequested,
authorize: workspaceAuthorizationMockFns.mockAuthorizeWorkspaceOperation,
authorize: workspaceAuthorizationMockFns.mockRequireCurrentHumanRole,
banned: authBanMockFns.mockGetActivelyBannedUserIds,
workspace: workspaceContextMockFns.mockResolveActiveWorkspaceApplicationContext,
}
Expand Down Expand Up @@ -60,7 +60,12 @@ beforeEach(() => {
allowPersonalApiKeys: false,
})
mocks.authorize.mockResolvedValue(undefined)
mocks.run.mockResolvedValue({ chatId, workspaceId: 'workspace' })
mocks.run.mockResolvedValue({
userId: 'actor',
chatId,
workspaceId: 'workspace',
organizationId: null,
})
mocks.stopped.mockResolvedValue(true)
})

Expand Down Expand Up @@ -103,15 +108,24 @@ describe('worker run control boundary', () => {
expect(response.status).toBe(200)
expect(await response.json()).toEqual({ stopped: false })
})
it.each([
null,
{ chatId: 'other-chat', workspaceId: 'workspace' },
{ chatId, workspaceId: 'other-workspace' },
])('answers Stop for a missing or mismatched canonical run: %j', async (run) => {
mocks.run.mockResolvedValue(run)
const response = await POST(request())
expect(response.status).toBe(200)
expect(await response.json()).toEqual({ stopped: true })
it.each([null, { chatId: 'other-chat', workspaceId: 'workspace' }])(
'answers Stop for a missing or mismatched canonical run: %j',
async (run) => {
mocks.run.mockResolvedValue(run)
const response = await POST(request())
expect(response.status).toBe(200)
expect(await response.json()).toEqual({ stopped: true })
expect(mocks.stopped).not.toHaveBeenCalled()
}
)
it('rejects caller scope that differs from original admission', async () => {
mocks.run.mockResolvedValue({
userId: 'actor',
chatId,
workspaceId: 'other-workspace',
organizationId: null,
})
expect((await POST(request())).status).toBe(403)
expect(mocks.stopped).not.toHaveBeenCalled()
})
it('does not turn a database outage into permission to continue', async () => {
Expand Down
2 changes: 1 addition & 1 deletion apps/sim/app/api/mothership/runs/control/route.ts
Original file line number Diff line number Diff line change
Expand Up @@ -13,7 +13,7 @@ import {

export const POST = defineInternalJsonRoute({
contract: readRunControlContract,
auth: internalCopilotAuth(RUN_CONTROL_AUDIENCE, { organization: true }),
auth: internalCopilotAuth(RUN_CONTROL_AUDIENCE, { organization: true, personal: true }),
operation: readRunControlOperation,
rateLimit: internalRateLimits.none({
reason: 'Worker reconciles durable cancellation for its own run.',
Expand Down
2 changes: 1 addition & 1 deletion apps/sim/app/api/mothership/tasks/workflow-status/route.ts
Original file line number Diff line number Diff line change
Expand Up @@ -11,7 +11,7 @@ import { readWatchedWorkflowStatus } from '@/lib/mothership/tasks/application/re

export const POST = defineInternalJsonRoute({
contract: workflowWatchStatusContract,
auth: internalCopilotAuth(TASK_DELEGATION_AUDIENCE, { organization: true }),
auth: internalCopilotAuth(TASK_DELEGATION_AUDIENCE, { organization: true, personal: true }),
operation: taskOperations.readWorkflowStatus,
rateLimit: internalRateLimits.none({
reason: 'Authenticated worker reconciles its bounded task inbox.',
Expand Down
2 changes: 1 addition & 1 deletion apps/sim/app/api/mothership/workspaces/context/route.ts
Original file line number Diff line number Diff line change
Expand Up @@ -13,7 +13,7 @@ import { WORKSPACE_TARGET_AUDIENCE } from '@/lib/mothership/chat/application/wor

export const POST = defineInternalJsonRoute({
contract: readWorkspaceContextContract,
auth: internalCopilotAuth(WORKSPACE_TARGET_AUDIENCE, { organization: true }),
auth: internalCopilotAuth(WORKSPACE_TARGET_AUDIENCE, { organization: true, personal: true }),
operation: readWorkspaceContextOperation,
rateLimit: internalRateLimits.none({
reason:
Expand Down
47 changes: 46 additions & 1 deletion apps/sim/lib/mothership/application/authorize-chat-callback.ts
Original file line number Diff line number Diff line change
Expand Up @@ -18,6 +18,8 @@ import {
COPILOT_VALIDATION_PURPOSE,
type CopilotValidationPurpose,
} from '@/lib/mothership/generated/billing-protocol-v1'
import { RunControlRequest } from '@/lib/mothership/generated/run-control'
import { authorizeOriginalRun } from '@/lib/mothership/request/application/run-authorization'
import { resolveActiveWorkspaceApplicationContext } from '@/lib/workspaces/application/context'

const CALLBACK_AUDIENCE = 'sim:copilot-callback'
Expand Down Expand Up @@ -48,6 +50,7 @@ const cancelWorkspaceChat = defineAuthorizedWorkspaceUseCase({
})

interface CopilotChatCallbackContext {
executionStreamId?: string
userId: string
workspaceId?: string
organizationId?: string
Expand All @@ -59,6 +62,44 @@ interface CopilotChatCallbackContext {

/** Reauthorizes the original server-owned scope across a model lifecycle boundary. */
export async function authorizeCopilotChatCallback(context: CopilotChatCallbackContext) {
if (context.executionStreamId !== undefined) {
if (
!context.chatId ||
(context.workspaceId && context.organizationId) ||
!RunControlRequest.shape.streamId.safeParse(context.executionStreamId).success
)
throw new OrchestrationError('forbidden', 'Invalid conversation scope')
const identity = { ...context, chatId: context.chatId }
const options = { audience: CALLBACK_AUDIENCE, ttlMs: COPILOT_APPLICATION_DELEGATION_TTL_MS }
const principal = context.organizationId
? createTrustedOrganizationCopilotPrincipal(
{ ...identity, organizationId: context.organizationId },
options
)
: context.workspaceId
? createTrustedCopilotPrincipal({ ...identity, workspaceId: context.workspaceId }, options)
: createTrustedCopilotPrincipal(
{
userId: identity.userId,
chatId: identity.chatId,
delegationId: identity.delegationId,
...(identity.executionStreamId
? { executionStreamId: identity.executionStreamId }
: {}),
},
options
)
await authorizeOriginalRun({
principal,
chatId: context.chatId,
streamId: context.executionStreamId,
audience: CALLBACK_AUDIENCE,
mode: context.mode,
capability:
context.purpose === COPILOT_VALIDATION_PURPOSE.cancellation ? 'none' : 'copilot.use',
})
return
}
if (context.organizationId) {
if (!context.chatId || context.workspaceId) {
throw new OrchestrationError('forbidden', 'Invalid conversation scope')
Expand All @@ -76,7 +117,11 @@ export async function authorizeCopilotChatCallback(context: CopilotChatCallbackC
await authorizeOrganizationChatDelegation.execute({ principal, mode: context.mode })
return
}
if (!context.workspaceId) return
if (!context.workspaceId)
throw new OrchestrationError(
'forbidden',
'Personal continuation requires original execution identity'
)

const principal = createTrustedCopilotPrincipal(
{ ...context, workspaceId: context.workspaceId },
Expand Down
30 changes: 26 additions & 4 deletions apps/sim/lib/mothership/application/workspace-target.ts
Original file line number Diff line number Diff line change
Expand Up @@ -3,6 +3,7 @@ import { OrchestrationError } from '@/lib/core/orchestration/types'
import {
COPILOT_APPLICATION_DELEGATION_TTL_MS,
createCopilotChatPrincipal,
createTrustedCopilotPrincipal,
createTrustedOrganizationCopilotPrincipal,
} from '@/lib/mothership/auth/application-delegation'
import {
Expand All @@ -12,6 +13,9 @@ import {
import { resolveActiveWorkspaceApplicationContext } from '@/lib/workspaces/application/context'

export interface InvocationOwner {
/** Server-owned execution correlation, never read from model CLI arguments. */
messageId?: string
copilotToolExecution?: boolean
userId: string
workspaceId?: string
organizationId?: string
Expand All @@ -31,20 +35,38 @@ export async function resolveInvocationWorkspace(owner: InvocationOwner, request
if (!organizationId && owner.workspaceId && workspaceId !== owner.workspaceId)
throw new OrchestrationError('not_found', 'Workspace not found in this conversation')
if (owner.chatId) {
const executionStreamId =
owner.copilotToolExecution && owner.messageId ? owner.messageId : undefined
const principal = organizationId
? createTrustedOrganizationCopilotPrincipal(
{
userId: owner.userId,
organizationId,
chatId: owner.chatId,
delegationId: `workspace-target:${owner.chatId}`,
...(executionStreamId ? { executionStreamId } : {}),
},
{ audience: WORKSPACE_TARGET_AUDIENCE, ttlMs: COPILOT_APPLICATION_DELEGATION_TTL_MS }
)
: createCopilotChatPrincipal(
{ userId: owner.userId, workspaceId, chatId: owner.chatId },
WORKSPACE_TARGET_AUDIENCE
)
: !owner.workspaceId
? createTrustedCopilotPrincipal(
{
userId: owner.userId,
chatId: owner.chatId,
delegationId: `workspace-target:${owner.chatId}`,
...(executionStreamId ? { executionStreamId } : {}),
},
{ audience: WORKSPACE_TARGET_AUDIENCE, ttlMs: COPILOT_APPLICATION_DELEGATION_TTL_MS }
)
: createCopilotChatPrincipal(
{
userId: owner.userId,
workspaceId,
chatId: owner.chatId,
...(executionStreamId ? { executionStreamId } : {}),
},
WORKSPACE_TARGET_AUDIENCE
)
return authorizeChatWorkspaceTarget.execute({
principal,
input: { chatId: owner.chatId, workspaceId },
Expand Down
16 changes: 16 additions & 0 deletions apps/sim/lib/mothership/async-runs/orphaned-runs.integration.ts
Original file line number Diff line number Diff line change
Expand Up @@ -34,6 +34,7 @@ import {
copilotAsyncToolCalls,
copilotChats,
copilotRequestStops,
copilotRunStops,
copilotRuns,
permissions,
user,
Expand Down Expand Up @@ -283,6 +284,21 @@ describe.runIf(Boolean(redisUrl))('Chat runs no controller owns', () => {
return row
}

it('settles personal runs from unified-only Stop intent', async () => {
const run = await admittedRun({ idleMinutes: 90 })
await db
.update(copilotRuns)
.set({ workspaceId: null, organizationId: null })
.where(eq(copilotRuns.id, run.runId))
await db.insert(copilotRunStops).values({ userId, streamId: run.streamId })
expect(await settleStoppedRunWithoutController(run.runId)).toBe(true)
const [settled] = await db
.select({ status: copilotRuns.status, error: copilotRuns.error })
.from(copilotRuns)
.where(eq(copilotRuns.id, run.runId))
expect(settled).toEqual({ status: 'cancelled', error: null })
})

it('leaves an admitted pending call unclaimed for its caller to settle as never run', async () => {
const run = await admittedRun()
const tool = await lateTool(run.runId)
Expand Down
7 changes: 6 additions & 1 deletion apps/sim/lib/mothership/async-runs/orphaned-runs.ts
Original file line number Diff line number Diff line change
Expand Up @@ -5,6 +5,7 @@ import {
copilotChats,
copilotOrganizationRequestStops,
copilotRequestStops,
copilotRunStops,
copilotRuns,
} from '@sim/db/schema'
import { createLogger } from '@sim/logger'
Expand Down Expand Up @@ -119,11 +120,15 @@ const legacyRunIdle = and(
const orphanIdle = or(leasedRunIdle, legacyRunIdle)

/** The user pressed Stop on this stream; a newer turn also closes tool admission, without one. */
const stopRequested = sql`(EXISTS (SELECT 1 FROM ${copilotRequestStops} s
const stopRequested = sql`(EXISTS (SELECT 1 FROM ${copilotRunStops} s
WHERE s.user_id = ${copilotRuns.userId} AND s.stream_id = ${copilotRuns.streamId})
OR EXISTS (SELECT 1 FROM ${copilotRequestStops} s
WHERE s.user_id = ${copilotRuns.userId} AND s.workspace_id = ${copilotRuns.workspaceId}
AND ${copilotRuns.organizationId} IS NULL
AND s.stream_id = ${copilotRuns.streamId})
OR EXISTS (SELECT 1 FROM ${copilotOrganizationRequestStops} s
WHERE s.user_id = ${copilotRuns.userId} AND s.organization_id = ${copilotRuns.organizationId}
AND ${copilotRuns.workspaceId} IS NULL
AND s.stream_id = ${copilotRuns.streamId}))`

export interface UnownedRun {
Expand Down
13 changes: 9 additions & 4 deletions apps/sim/lib/mothership/async-runs/repository.test.ts
Loading
Loading