fix(mothership): let the agent copy files between workspaces without Projects by waleedlatif1 · Pull Request #9288 · simstudioai/sim · GitHub
Skip to content
Merged
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
35 changes: 21 additions & 14 deletions apps/sim/lib/mothership/application/file-owner-context.ts
2 changes: 1 addition & 1 deletion apps/sim/lib/mothership/chat/post.ts
Original file line number Diff line number Diff line change
Expand Up @@ -1408,7 +1408,7 @@ export async function handleUnifiedChatPost(req: NextRequest) {
if (!lease || !sendClaim?.claimToken || sendClaim.storageMethod !== 'database') {
throw new Error('Chat admission requires its controller and durable send claim')
}
const resourceAuthoringVersion = await durableResourceAuthoringVersion(
const resourceAuthoringVersion = durableResourceAuthoringVersion(
executionContext.copilotResourceAdmission !== undefined
)
admittedRun = await admitChatTurn.execute({
Expand Down
10 changes: 9 additions & 1 deletion apps/sim/lib/mothership/generated/protocol.ts
Original file line number Diff line number Diff line change
Expand Up @@ -303,6 +303,13 @@ export const TOOL_DEADLINE_SETTLE_MARGIN_MS = 90_000;
/** The `clientCapabilities` token a Sim sends when it honors `StreamToolDeadline`. */
export const TOOL_DEADLINE_CAPABILITY = "tool-deadline";

/**
* The `clientCapabilities` token a Sim sends when this turn may copy files within or between
* workspaces it can reach (`files copy`) without negotiated file owners, so also where Projects
* are off.
*/
export const WORKSPACE_FILE_COPY_CAPABILITY = "workspace-file-copy";

/** POST /api/mothership — the chat request sim sends. */
export interface ChatRequest extends StreamResponseReceipt {
fileOwnerProtocolVersion?: 1 | undefined;
Expand Down Expand Up @@ -348,7 +355,8 @@ export interface ChatRequest extends StreamResponseReceipt {
* grace windows and run tools server-side immediately. ABSENT = older/unknown caller —
* dispatch keeps its conservative grace (deploy-skew safe: a stale tab that predates
* this field still gets waited on). Known capability: "workflow-tool-pickup". Sim also
* declares its own {@link TOOL_DEADLINE_CAPABILITY} here.
* declares its own {@link TOOL_DEADLINE_CAPABILITY} and {@link WORKSPACE_FILE_COPY_CAPABILITY}
* here.
*/
clientCapabilities?: string[] | undefined;
/**
Expand Down
25 changes: 16 additions & 9 deletions apps/sim/lib/mothership/request/lifecycle/recovery-config.test.ts
Original file line number Diff line number Diff line change
Expand Up @@ -104,15 +104,22 @@ describe('recovery marker written by an admitted turn', () => {
}),
}

it('stays off a recovery config older builds must still parse while Projects are off', async () => {
featureFlagsMockFns.mockIsFeatureEnabled.mockResolvedValue(false)
const recovery = await streamRecoveryConfig(options, workspaceRequest)
expect(recovery).not.toHaveProperty('resourceAuthoringVersion')
})
it.each([true, false])(
'restores authoring admission on recovery whether Projects are on (%s) or off',
async (projects) => {
featureFlagsMockFns.mockIsFeatureEnabled.mockImplementation(
async (flag) => projects && flag === 'projects'
)
const recovery = await streamRecoveryConfig(options, workspaceRequest)
expect(recovery?.resourceAuthoringVersion).toBe(1)
}
)

it('restores authoring admission on recovery while Projects are on', async () => {
featureFlagsMockFns.mockIsFeatureEnabled.mockImplementation(async (flag) => flag === 'projects')
const recovery = await streamRecoveryConfig(options, workspaceRequest)
expect(recovery?.resourceAuthoringVersion).toBe(1)
it('stays off a turn that was never admitted', async () => {
const recovery = await streamRecoveryConfig(
{ ...options, copilotResourceAdmission: undefined },
workspaceRequest
)
expect(recovery).not.toHaveProperty('resourceAuthoringVersion')
})
})
15 changes: 7 additions & 8 deletions apps/sim/lib/mothership/request/lifecycle/recovery-config.ts
Original file line number Diff line number Diff line change
Expand Up @@ -3,7 +3,6 @@ import { hasCopilotResourceAdmission } from '@/lib/mothership/auth/application-d
import { COPILOT_REQUEST_MODES } from '@/lib/mothership/constants'
import { ChatPayloadSchema } from '@/lib/mothership/generated/protocol'
import type { CopilotLifecycleOptions } from '@/lib/mothership/request/lifecycle/run'
import { isProjectsEnabled } from '@/lib/projects/rollout.server'

export const BillingAdmissionSchema = z
.object({
Expand All @@ -28,14 +27,14 @@ export const DurableChatRequestSchema = ChatPayloadSchema.safeExtend({
export const RESOURCE_AUTHORING_VERSION = 1

/**
* The marker an admitted turn stores for recovery. Builds before it parse recovery strictly and
* would refuse to recover the run, and only Project authoring reads it, so it is written only
* while Projects are on.
* The marker an admitted turn stores for recovery, so a recovered turn keeps its admission. It is
* written whether or not Projects are on, since copies between workspaces need the admission
* without them; every release since the marker shipped parses it.
*/
export async function durableResourceAuthoringVersion(
export function durableResourceAuthoringVersion(
admitted: boolean
): Promise<typeof RESOURCE_AUTHORING_VERSION | undefined> {
return admitted && (await isProjectsEnabled()) ? RESOURCE_AUTHORING_VERSION : undefined
): typeof RESOURCE_AUTHORING_VERSION | undefined {
return admitted ? RESOURCE_AUTHORING_VERSION : undefined
}

export const StreamRecoveryConfigSchema = z
Expand Down Expand Up @@ -66,7 +65,7 @@ export async function streamRecoveryConfig(
request: Record<string, unknown>
) {
if (options.interactive !== true) return undefined
const resourceAuthoringVersion = await durableResourceAuthoringVersion(
const resourceAuthoringVersion = durableResourceAuthoringVersion(
hasCopilotResourceAdmission({
...options,
...options.executionContext,
Expand Down
16 changes: 11 additions & 5 deletions apps/sim/lib/mothership/tools/sandbox-file-copy.ts
Original file line number Diff line number Diff line change
Expand Up @@ -13,6 +13,7 @@ import {
recordSandboxResourceEffects,
type SandboxResourceScope,
} from '@/lib/mothership/tools/sandbox-resources'
import { copiesWithinWorkspaces } from '@/lib/workspace-files/application/copy-operation'

const logger = createLogger('SandboxFileCopy')

Expand All @@ -22,15 +23,17 @@ export async function proxySandboxFileCopyRequest(
token: string,
scope: SandboxResourceScope
): Promise<Response> {
const unavailable = () =>
Response.json(
{ error: 'Paired file owner routing is unavailable or contradictory' },
{ status: 403 }
)
if (
scope.fileOwnerProtocolVersion !== 1 ||
(scope.fileOwnerProtocolVersion !== 1 && !scope.workspaceFileCopy) ||
request.headers.has('x-mothership-file-owner') ||
request.headers.has('x-mothership-workspace-id')
)
return Response.json(
{ error: 'Paired file owner routing is unavailable or contradictory' },
{ status: 403 }
)
return unavailable()
if (request.method !== 'POST')
return Response.json({ error: 'Method not allowed' }, { status: 405 })
const parsed = await parseRequest(
Expand All @@ -43,6 +46,9 @@ export async function proxySandboxFileCopyRequest(
}
)
if (!parsed.success) return parsed.response
// Without the negotiated protocol only workspace owners exist for this lease.
if (scope.fileOwnerProtocolVersion !== 1 && !copiesWithinWorkspaces(parsed.data.body))
return unavailable()
const endpoint = getInternalApiBaseUrl()
const transport = createFileCopyCliTransport(
endpoint,
Expand Down
65 changes: 56 additions & 9 deletions apps/sim/lib/mothership/tools/sandbox-resource-transport.test.ts
Original file line number Diff line number Diff line change
Expand Up @@ -16,6 +16,7 @@ import { assertWorkspaceInvocationScope } from '@/lib/core/application/workspace
import { OrchestrationError } from '@/lib/core/orchestration/types'
import { chatSandboxSessionKey } from '@/lib/mothership/tools/sandbox-session-key'
import { projectFileOperations } from '@/lib/projects/files/application/operations'
import { fileCopyOperation } from '@/lib/workspace-files/application/copy-operation'
import { fileOperations } from '@/lib/workspace-files/application/operations'

const { readScope, recordEffects, fetcher, routeMatcher, recordInput } = vi.hoisted(() => ({
Expand Down Expand Up @@ -436,26 +437,72 @@ it.each([
}
)

it('refuses a Project copy destination on a lease without the file-owner protocol', async () => {
readScope.mockResolvedValue({ ...scope, fileOwnerProtocolVersion: undefined })
routeMatcher.mockReturnValue({ params: {}, load: async () => ({ POST: fetcher }) })
it('binds a workspace pair under the paired copy grant on an admitted lease without the protocol', async () => {
readScope.mockResolvedValue({
...scope,
fileOwnerProtocolVersion: undefined,
workspaceFileCopy: true,
})
routeMatcher.mockReturnValue({
pattern: '/api/v2/files/copy',
params: {},
load: async () => ({ POST: fetcher }),
})
fetcher.mockImplementation(async (forwarded: Request) => {
expect(
copilotRequestPrincipal(forwarded, fileCopyOperation, {
operation: fileCopyOperation,
delegationAudience: fileCopyOperation.delegationAudience,
})
).toMatchObject({
kind: 'resource_delegated',
subjectUserId: scope.userId,
invocation: { kind: 'chat', chatId: scope.chatId },
scope: { kind: 'file_copy', destination: { owner: workspaceOwner } },
})
return Response.json({ data: { copied: [] } })
})
const response = await proxySandboxResourceRequest(
request('/api/v2/files/copy', {
method: 'POST',
body: JSON.stringify({
source: { owner: workspaceOwner, fileIds: ['file'] },
destination: { owner: { entityType: 'project', entityId: 'project' } },
destination: { owner: workspaceOwner },
}),
}),
token
)
expect(response.status).toBe(403)
expect(await response.json()).toEqual({
error: 'Paired file owner routing is unavailable or contradictory',
})
expect(fetcher).not.toHaveBeenCalled()
expect(response.status).toBe(200)
expect(fetcher).toHaveBeenCalledTimes(1)
})

it.each([undefined, true as const])(
'refuses a Project copy destination on a lease without the file-owner protocol (workspace copies %s)',
async (workspaceFileCopy) => {
readScope.mockResolvedValue({
...scope,
fileOwnerProtocolVersion: undefined,
workspaceFileCopy,
})
routeMatcher.mockReturnValue({ params: {}, load: async () => ({ POST: fetcher }) })
const response = await proxySandboxResourceRequest(
request('/api/v2/files/copy', {
method: 'POST',
body: JSON.stringify({
source: { owner: workspaceOwner, fileIds: ['file'] },
destination: { owner: { entityType: 'project', entityId: 'project' } },
}),
}),
token
)
expect(response.status).toBe(403)
expect(await response.json()).toEqual({
error: 'Paired file owner routing is unavailable or contradictory',
})
expect(fetcher).not.toHaveBeenCalled()
}
)

it('validates a copy body itself once the lease negotiated the file-owner protocol', async () => {
readScope.mockResolvedValue({ ...scope, fileOwnerProtocolVersion: 1 })
routeMatcher.mockReturnValue({ params: {}, load: async () => ({ POST: fetcher }) })
Expand Down
18 changes: 10 additions & 8 deletions apps/sim/lib/mothership/tools/sandbox-resource-transport.ts
Original file line number Diff line number Diff line change
Expand Up @@ -28,6 +28,7 @@ import {
type SandboxResourceScope,
} from '@/lib/mothership/tools/sandbox-resources'
import { chatSandboxSessionKey } from '@/lib/mothership/tools/sandbox-session-key'
import { copiesWithinWorkspaces } from '@/lib/workspace-files/application/copy-operation'

const logger = createLogger('MothershipSandboxResourceTransport')

Expand All @@ -54,15 +55,20 @@ export async function proxySandboxResourceRequest(
if (!scope)
return Response.json({ error: 'Sandbox tool execution is no longer active' }, { status: 403 })
/**
* Paired owners need the negotiated protocol. Without it, only a copy that names a non-workspace
* owner is refused there; workspace copies and invalid bodies keep the workspace route and its
* Copies bind both owners under the paired grant. Project owners need the negotiated protocol,
* and an admitted lease copies between workspaces without it; a copy naming a Project owner is
* refused there. Otherwise workspace copies and invalid bodies keep the workspace route and its
* v2 validation response.
*/
if (path === v2CopyFileItemsContract.path) {
const buffered = await bufferCopyRequest(request)
if (buffered instanceof Response) return buffered
request = buffered
if (scope.fileOwnerProtocolVersion === 1 || (await copiesBeyondWorkspaces(request)))
if (
Comment thread
waleedlatif1 marked this conversation as resolved.
scope.fileOwnerProtocolVersion === 1 ||
scope.workspaceFileCopy ||
(await copiesBeyondWorkspaces(request))
)
return proxySandboxFileCopyRequest(request, token, scope)
}
const ownerHeader = request.headers.get('x-mothership-file-owner')
Expand Down Expand Up @@ -245,9 +251,5 @@ async function copiesBeyondWorkspaces(request: Request): Promise<boolean> {
{},
V2_PARSE_DEFAULTS
)
return (
parsed.success &&
(parsed.data.body.source.owner.entityType !== 'workspace' ||
parsed.data.body.destination.owner.entityType !== 'workspace')
)
return parsed.success && !copiesWithinWorkspaces(parsed.data.body)
}
3 changes: 2 additions & 1 deletion apps/sim/lib/mothership/tools/sandbox-resources.test.ts
Loading
Loading