Skip to content

Commit 68c4dc9

Browse files
Sg312Sg312
committed
perf(mothership): trace CLI and tool persistence boundaries
1 parent c05c34e commit 68c4dc9

12 files changed

Lines changed: 522 additions & 133 deletions

File tree

‎apps/sim/lib/mothership/agent-cli/file-read-transport.ts‎

Lines changed: 3 additions & 2 deletions
Original file line numberDiff line numberDiff line change
@@ -27,6 +27,7 @@ const logger = createLogger('MothershipFileReads')
2727
/** Private file-read metadata stays in the authenticated host; the CLI consumes its usual wire shape. */
2828
export function createFileReadTransport(context: {
2929
endpoint: string
30+
transport?: typeof fetch
3031
userId: string
3132
chatId?: string
3233
registry?: ResolvedSecretTraceRegistry
@@ -63,10 +64,10 @@ export function createFileReadTransport(context: {
6364
!url.pathname.startsWith(prefix) ||
6465
collectionPaths.has(url.pathname)
6566
) {
66-
return fetch(input, init)
67+
return (context.transport ?? fetch)(input, init)
6768
}
6869
const match = /^([^/]+)(\/text)?$/.exec(url.pathname.slice(prefix.length))
69-
if (!match) return fetch(input, init)
70+
if (!match) return (context.transport ?? fetch)(input, init)
7071
const request = new NextRequest(new Request(input, init))
7172
let stream: ReadableStream<Uint8Array> | undefined
7273
try {

‎apps/sim/lib/mothership/agent-cli/index.ts‎

Lines changed: 39 additions & 23 deletions
Original file line numberDiff line numberDiff line change
@@ -10,10 +10,13 @@ import { createFileReadTransport } from '@/lib/mothership/agent-cli/file-read-tr
1010
import { createFileUploadTransport } from '@/lib/mothership/agent-cli/file-upload-transport'
1111
import { runCli } from '@/lib/mothership/agent-cli/run-cli'
1212
import { applySink } from '@/lib/mothership/agent-cli/sink'
13+
import { createTracedCliTransport } from '@/lib/mothership/agent-cli/traced-transport'
1314
import { agentCliFail } from '@/lib/mothership/agent-cli/types'
1415
import { createWorkbenchFileProvenance } from '@/lib/mothership/agent-cli/workbench-file-provenance'
1516
import { mintDelegationToken } from '@/lib/mothership/chat/delegation'
1617
import type { AgentCliRawResult, AgentCliRequest } from '@/lib/mothership/generated/agent-cli'
18+
import { TraceSpan } from '@/lib/mothership/generated/trace-spans-v1'
19+
import { withCopilotSpan } from '@/lib/mothership/request/otel'
1720
import { chatSandboxSessionKey } from '@/lib/mothership/tools/sandbox-session-key'
1821
import type { ResolvedSecretTraceRegistry } from '@/executor/utils/resolved-secret-trace-registry'
1922

@@ -36,16 +39,19 @@ export async function executeAgentCliRequest(
3639
context: AgentCliExecutionContext
3740
): Promise<AgentCliRawResult> {
3841
context.signal?.throwIfAborted()
39-
const apiKey = await mintDelegationToken({
40-
workspaceId: context.workspaceId,
41-
userId: context.userId,
42-
})
42+
const apiKey = await withCopilotSpan(TraceSpan.CopilotCliIdentity, undefined, () =>
43+
mintDelegationToken({
44+
workspaceId: context.workspaceId,
45+
userId: context.userId,
46+
})
47+
)
4348
if (!apiKey) return agentCliFail('Could not establish workspace credentials for this command.')
4449
const endpoint = getInternalApiBaseUrl()
4550
const sessionKey = context.chatId ? chatSandboxSessionKey(context.chatId) : null
4651
const files = sessionKey ? createWorkbenchFileProvenance({ ...context, sessionKey }) : undefined
4752
const reads = createFileReadTransport({
4853
endpoint,
54+
transport: createTracedCliTransport(endpoint, fetch),
4955
userId: context.userId,
5056
registry: context.resolvedSecretTraceRegistry,
5157
...(context.chatId !== undefined ? { chatId: context.chatId } : {}),
@@ -67,33 +73,43 @@ export async function executeAgentCliRequest(
6773
...(context.signal ? { signal: context.signal } : {}),
6874
}
6975

76+
const { invocation, sink } = request
7077
let result: AgentCliRawResult
7178
context.signal?.throwIfAborted()
72-
if (request.invocation.kind === 'stdout') {
79+
if (invocation.kind === 'stdout') {
7380
// Text the worker already holds (sliced, or worker-answered): only the sink applies.
74-
result = { exitCode: 0, stdout: request.invocation.stdout, stderr: '' }
75-
} else if (request.invocation.kind === 'augmentation') {
76-
result = await runEngine(
77-
request.invocation.name,
78-
request.invocation.positionals,
79-
{
80-
client: createEmbeddedClient(identity),
81-
workspaceId: context.workspaceId,
82-
userId: context.userId,
83-
principal: await principalForDelegation(apiKey),
84-
...(context.chatId !== undefined ? { chatId: context.chatId } : {}),
85-
signal: context.signal,
86-
},
87-
request.invocation.flags
81+
result = { exitCode: 0, stdout: invocation.stdout, stderr: '' }
82+
} else if (invocation.kind === 'augmentation') {
83+
result = await withCopilotSpan(TraceSpan.CopilotCliInvoke, undefined, async () =>
84+
runEngine(
85+
invocation.name,
86+
invocation.positionals,
87+
{
88+
client: createEmbeddedClient(identity),
89+
workspaceId: context.workspaceId,
90+
userId: context.userId,
91+
principal: await principalForDelegation(apiKey),
92+
...(context.chatId !== undefined ? { chatId: context.chatId } : {}),
93+
signal: context.signal,
94+
},
95+
invocation.flags
96+
)
8897
)
8998
} else {
90-
result = await runCli(request.invocation.argv, identity, sessionKey, files)
99+
const { argv } = invocation
100+
result = await withCopilotSpan(TraceSpan.CopilotCliInvoke, undefined, () =>
101+
runCli(argv, identity, sessionKey, files)
102+
)
91103
if (result.exitCode === 0 && request.curate === 'block') {
92-
result = await curateBlockDetail(result, context)
104+
result = await withCopilotSpan(TraceSpan.CopilotCliCurate, undefined, () =>
105+
curateBlockDetail(result, context)
106+
)
93107
}
94108
}
95-
return request.sink
96-
? applySink(request.sink, sessionKey, result, context.signal, files?.observeOutput)
109+
return sink
110+
? withCopilotSpan(TraceSpan.CopilotCliSink, undefined, () =>
111+
applySink(sink, sessionKey, result, context.signal, files?.observeOutput)
112+
)
97113
: result
98114
}
99115

‎apps/sim/lib/mothership/agent-cli/traced-transport.ts‎

Lines changed: 39 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -0,0 +1,39 @@
1+
import { SpanKind, SpanStatusCode } from '@opentelemetry/api'
2+
import { TraceAttr } from '@/lib/mothership/generated/trace-attributes-v1'
3+
import { TraceSpan } from '@/lib/mothership/generated/trace-spans-v1'
4+
import { traceHeaders } from '@/lib/mothership/request/go/propagation'
5+
import { getCopilotTracer } from '@/lib/mothership/request/otel'
6+
7+
/** Links internal CLI HTTP calls to their tool; the span ends at response headers. */
8+
export function createTracedCliTransport(endpoint: string, transport: typeof fetch): typeof fetch {
9+
const origin = new URL(endpoint).origin
10+
return async (input, init) => {
11+
const url = new URL(input instanceof Request ? input.url : input)
12+
if (url.origin !== origin) return transport(input, init)
13+
const method = init?.method ?? (input instanceof Request ? input.method : 'GET')
14+
return getCopilotTracer().startActiveSpan(
15+
TraceSpan.CopilotCliHttpHeaders,
16+
{
17+
kind: SpanKind.CLIENT,
18+
attributes: { [TraceAttr.HttpMethod]: method, [TraceAttr.HttpPath]: url.pathname },
19+
},
20+
async (span) => {
21+
try {
22+
const headers = new Headers(
23+
init?.headers ?? (input instanceof Request ? input.headers : undefined)
24+
)
25+
for (const [key, value] of Object.entries(traceHeaders())) headers.set(key, value)
26+
const response = await transport(input, { ...init, headers })
27+
span.setAttribute(TraceAttr.HttpStatusCode, response.status)
28+
if (!response.ok) span.setStatus({ code: SpanStatusCode.ERROR })
29+
return response
30+
} catch (error) {
31+
span.setStatus({ code: SpanStatusCode.ERROR })
32+
throw error
33+
} finally {
34+
span.end()
35+
}
36+
}
37+
)
38+
}
39+
}

‎apps/sim/lib/mothership/async-runs/repository.ts‎

Lines changed: 88 additions & 70 deletions
Original file line numberDiff line numberDiff line change
@@ -1,4 +1,4 @@
1-
import { trace } from '@opentelemetry/api'
1+
import { SpanKind, trace } from '@opentelemetry/api'
22
import { db } from '@sim/db'
33
import {
44
type CopilotAsyncToolStatus,
@@ -34,6 +34,10 @@ import {
3434
} from '@/lib/mothership/async-runs/execution-lease'
3535
import { TraceAttr } from '@/lib/mothership/generated/trace-attributes-v1'
3636
import { TraceSpan } from '@/lib/mothership/generated/trace-spans-v1'
37+
import {
38+
traceMothershipQuery,
39+
traceMothershipTransaction,
40+
} from '@/lib/mothership/observability/database'
3741
import { markSpanForError } from '@/lib/mothership/request/otel'
3842
import { chatSandboxSessionKey } from '@/lib/mothership/tools/sandbox-session-key'
3943
import {
@@ -71,22 +75,28 @@ async function withDbSpan<T>(
7175
attrs: Record<string, string | number | boolean | undefined>,
7276
fn: () => Promise<T>
7377
): Promise<T> {
74-
const span = getAsyncRunsTracer().startSpan(name, {
75-
attributes: {
76-
[TraceAttr.DbSystem]: 'postgresql',
77-
[TraceAttr.DbOperation]: op,
78-
[TraceAttr.DbSqlTable]: table,
79-
...filterUndefined(attrs),
78+
return getAsyncRunsTracer().startActiveSpan(
79+
name,
80+
{
81+
kind: SpanKind.CLIENT,
82+
attributes: {
83+
[TraceAttr.DbSystem]: 'postgresql',
84+
[TraceAttr.DbOperation]: op,
85+
[TraceAttr.DbSqlTable]: table,
86+
...filterUndefined(attrs),
87+
},
8088
},
81-
})
82-
try {
83-
return await fn()
84-
} catch (error) {
85-
markSpanForError(span, error)
86-
throw error
87-
} finally {
88-
span.end()
89-
}
89+
async (span) => {
90+
try {
91+
return await fn()
92+
} catch (error) {
93+
markSpanForError(span, error)
94+
throw error
95+
} finally {
96+
span.end()
97+
}
98+
}
99+
)
90100
}
91101

92102
export interface CreateRunSegmentInput {
@@ -366,20 +376,22 @@ export async function upsertAsyncToolCall(input: {
366376
const now = new Date()
367377
const args = sanitizeValueForJsonb(input.args ?? {})
368378
const sealedContext = sanitizeValueForJsonb(input.sealedContext)
369-
const [row] = await db
370-
.insert(copilotAsyncToolCalls)
371-
.values({
372-
runId: effectiveRunId,
373-
checkpointId: input.checkpointId ?? null,
374-
toolCallId: input.toolCallId,
375-
toolName: input.toolName,
376-
args,
377-
status: incomingStatus,
378-
...(sealedContext !== undefined ? { result: sealedContext } : {}),
379-
updatedAt: now,
380-
})
381-
.onConflictDoNothing()
382-
.returning()
379+
const [row] = await traceMothershipQuery('INSERT', 'copilot_async_tool_calls', () =>
380+
db
381+
.insert(copilotAsyncToolCalls)
382+
.values({
383+
runId: effectiveRunId,
384+
checkpointId: input.checkpointId ?? null,
385+
toolCallId: input.toolCallId,
386+
toolName: input.toolName,
387+
args,
388+
status: incomingStatus,
389+
...(sealedContext !== undefined ? { result: sealedContext } : {}),
390+
updatedAt: now,
391+
})
392+
.onConflictDoNothing()
393+
.returning()
394+
)
383395

384396
return row ?? getAsyncToolCall(input.toolCallId)
385397
}
@@ -491,54 +503,60 @@ export async function claimSimToolExecution(
491503
[TraceAttr.RunId]: input.runId,
492504
},
493505
() =>
494-
db.transaction(async (tx) => {
495-
const [run] = await tx
496-
.select({
497-
toolExecutionVersion: copilotRuns.toolExecutionVersion,
498-
toolAdmissionClosedAt: copilotRuns.toolAdmissionClosedAt,
499-
status: copilotRuns.status,
500-
})
501-
.from(copilotRuns)
502-
.where(and(eq(copilotRuns.id, input.runId), eq(copilotRuns.userId, input.userId)))
503-
.for('update')
506+
traceMothershipTransaction<SimToolExecutionClaim>('claim_tool', async (tx) => {
507+
const [run] = await traceMothershipQuery('SELECT FOR UPDATE', 'copilot_runs', () =>
508+
tx
509+
.select({
510+
toolExecutionVersion: copilotRuns.toolExecutionVersion,
511+
toolAdmissionClosedAt: copilotRuns.toolAdmissionClosedAt,
512+
status: copilotRuns.status,
513+
})
514+
.from(copilotRuns)
515+
.where(and(eq(copilotRuns.id, input.runId), eq(copilotRuns.userId, input.userId)))
516+
.for('update')
517+
)
504518
if (!run || run.toolExecutionVersion !== SIM_TOOL_EXECUTION_VERSION)
505519
throw new Error('Tool execution ownership is unavailable for this run')
506520
if (run.toolAdmissionClosedAt || TERMINAL_RUN_STATUSES.includes(run.status))
507521
return { outcome: 'closed' }
508522
const startedAt = new Date()
509-
const [claimed] = await tx
510-
.update(copilotAsyncToolCalls)
511-
.set({
512-
status: ASYNC_TOOL_STATUS.running,
513-
claimedBy: 'sim-stream',
514-
claimedAt: startedAt,
515-
executionStartedAt: startedAt,
516-
executionOwnerToken: input.ownerToken,
517-
executionLeaseExpiresAt: sql`clock_timestamp() + ${SIM_TOOL_EXECUTION_LEASE_SECONDS} * interval '1 second'`,
518-
updatedAt: startedAt,
519-
})
520-
.where(
521-
and(
522-
eq(copilotAsyncToolCalls.toolCallId, input.toolCallId),
523-
eq(copilotAsyncToolCalls.runId, input.runId),
524-
isNull(copilotAsyncToolCalls.executionStartedAt),
525-
inArray(copilotAsyncToolCalls.status, [
526-
ASYNC_TOOL_STATUS.pending,
527-
ASYNC_TOOL_STATUS.running,
528-
])
523+
const [claimed] = await traceMothershipQuery('UPDATE', 'copilot_async_tool_calls', () =>
524+
tx
525+
.update(copilotAsyncToolCalls)
526+
.set({
527+
status: ASYNC_TOOL_STATUS.running,
528+
claimedBy: 'sim-stream',
529+
claimedAt: startedAt,
530+
executionStartedAt: startedAt,
531+
executionOwnerToken: input.ownerToken,
532+
executionLeaseExpiresAt: sql`clock_timestamp() + ${SIM_TOOL_EXECUTION_LEASE_SECONDS} * interval '1 second'`,
533+
updatedAt: startedAt,
534+
})
535+
.where(
536+
and(
537+
eq(copilotAsyncToolCalls.toolCallId, input.toolCallId),
538+
eq(copilotAsyncToolCalls.runId, input.runId),
539+
isNull(copilotAsyncToolCalls.executionStartedAt),
540+
inArray(copilotAsyncToolCalls.status, [
541+
ASYNC_TOOL_STATUS.pending,
542+
ASYNC_TOOL_STATUS.running,
543+
])
544+
)
529545
)
530-
)
531-
.returning({ id: copilotAsyncToolCalls.id })
546+
.returning({ id: copilotAsyncToolCalls.id })
547+
)
532548
if (claimed) return { outcome: 'claimed' }
533-
const [record] = await tx
534-
.select({ id: copilotAsyncToolCalls.id })
535-
.from(copilotAsyncToolCalls)
536-
.where(
537-
and(
538-
eq(copilotAsyncToolCalls.toolCallId, input.toolCallId),
539-
eq(copilotAsyncToolCalls.runId, input.runId)
549+
const [record] = await traceMothershipQuery('SELECT', 'copilot_async_tool_calls', () =>
550+
tx
551+
.select({ id: copilotAsyncToolCalls.id })
552+
.from(copilotAsyncToolCalls)
553+
.where(
554+
and(
555+
eq(copilotAsyncToolCalls.toolCallId, input.toolCallId),
556+
eq(copilotAsyncToolCalls.runId, input.runId)
557+
)
540558
)
541-
)
559+
)
542560
if (!record) throw new Error('Tool execution record is unavailable')
543561
return { outcome: 'existing' }
544562
})

0 commit comments

Comments
 (0)