fix(mothership): keep display activity out of integration and mcp tool arguments by waleedlatif1 · Pull Request #9284 · 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
28 changes: 28 additions & 0 deletions apps/sim/lib/mothership/chat/persisted-message.test.ts
16 changes: 11 additions & 5 deletions apps/sim/lib/mothership/chat/persisted-message.ts
Original file line number Diff line number Diff line change
Expand Up @@ -24,6 +24,7 @@ import {
MothershipStreamV1ToolOutcome,
MothershipStreamV1ToolPhase,
} from '@/lib/mothership/generated/mothership-stream-v1'
import type { ToolActivity } from '@/lib/mothership/generated/protocol'
import type {
ContentBlock,
LocalToolCallStatus,
Expand Down Expand Up @@ -273,6 +274,14 @@ export function displayContentOf(content: PersistedMessage): PersistedMessage |

// Write: OrchestratorResult → PersistedMessage

/** A call's arguments as history reads them, with an inherited display grouping restored. */
function persistedParams(
params: Record<string, unknown> | undefined,
activity: ToolActivity | undefined
): Record<string, unknown> | undefined {
return activity && params?.activity === undefined ? { ...params, activity } : params
}

function resolveToolState(block: ContentBlock): PersistedToolState {
const tc = block.toolCall
if (!tc) return 'pending'
Expand Down Expand Up @@ -382,17 +391,14 @@ function mapContentBlockBody(block: ContentBlock): PersistedContentBlock {
const activityDescription = normalizeToolActivityDescription(
block.toolCall.activityDescription
)
const params = persistedParams(block.toolCall.params, block.toolCall.activity)
const toolCall: PersistedToolCall = {
id: block.toolCall.id,
name: block.toolCall.name,
state,
...(activityDescription ? { activityDescription } : {}),
...(isSubagentTool && isNonTerminal ? {} : { result: redactedResult }),
...(isSubagentTool && isNonTerminal
? {}
: block.toolCall.params
? { params: block.toolCall.params }
: {}),
...(isSubagentTool && isNonTerminal ? {} : params ? { params } : {}),
...(block.calledBy ? { calledBy: block.calledBy } : {}),
...(block.toolCall.displayTitle
? {
Expand Down
107 changes: 107 additions & 0 deletions apps/sim/lib/mothership/request/handlers/handlers.test.ts
Original file line number Diff line number Diff line change
@@ -1,3 +1,4 @@
import { flushMacrotask } from '@sim/testing/helpers/async'
import { createDeferred } from '@sim/testing/helpers/deferred'
import {
mothershipAsyncRunsMock,
Expand Down Expand Up @@ -2252,6 +2253,112 @@ describe('sse-handlers tool lifecycle', () => {
)
})

it("executes a resolved MCP operation with its own arguments under the gateway's grouping", async () => {
isSimExecuted.mockReturnValue(false)
executeTool.mockResolvedValueOnce({ success: true, output: {} })
const grouping = { id: 'gauge-read', title: 'Reading Gauge', completedTitle: 'Read Gauge' }

await sseHandlers.tool(
{
type: MothershipStreamV1EventType.tool,
payload: {
toolCallId: 'gateway-mcp',
toolName: 'call_integration_tool',
arguments: {
toolId: 'mcp-server-get_overview',
description: 'Pulling the overview',
activity: grouping,
arguments: { category_id: null },
},
executor: MothershipStreamV1ToolExecutor.go,
mode: MothershipStreamV1ToolMode.sync,
phase: MothershipStreamV1ToolPhase.call,
},
} satisfies StreamEvent,
context,
execContext,
{ interactive: false, timeout: 1000 }
)

await sseHandlers.tool(
{
type: MothershipStreamV1EventType.tool,
payload: {
toolCallId: 'gateway-mcp',
toolName: 'mcp-server-get_overview',
arguments: { category_id: null },
executor: MothershipStreamV1ToolExecutor.sim,
mode: MothershipStreamV1ToolMode.async,
phase: MothershipStreamV1ToolPhase.call,
},
} satisfies StreamEvent,
context,
execContext,
{ interactive: false, timeout: 1000 }
)

await flushMacrotask()

expect(executeTool).toHaveBeenCalledWith(
'mcp-server-get_overview',
{ category_id: null },
expect.any(Object)
)
expect(context.toolCalls.get('gateway-mcp')?.activity).toEqual(grouping)
})

it("keeps an integration operation's own activity argument under the gateway's grouping", async () => {
isSimExecuted.mockReturnValue(false)
executeTool.mockResolvedValueOnce({ success: true, output: {} })

await sseHandlers.tool(
{
type: MothershipStreamV1EventType.tool,
payload: {
toolCallId: 'gateway-work-item',
toolName: 'call_integration_tool',
arguments: {
toolId: 'azure_devops_create_work_item',
description: 'Filing the task',
activity: { id: 'filing', title: 'Filing tasks', completedTitle: 'Filed tasks' },
arguments: { title: 'Ship it', activity: 'Development' },
},
executor: MothershipStreamV1ToolExecutor.go,
mode: MothershipStreamV1ToolMode.sync,
phase: MothershipStreamV1ToolPhase.call,
},
} satisfies StreamEvent,
context,
execContext,
{ interactive: false, timeout: 1000 }
)

await sseHandlers.tool(
{
type: MothershipStreamV1EventType.tool,
payload: {
toolCallId: 'gateway-work-item',
toolName: 'azure_devops_create_work_item',
arguments: { title: 'Ship it', activity: 'Development' },
executor: MothershipStreamV1ToolExecutor.sim,
mode: MothershipStreamV1ToolMode.async,
phase: MothershipStreamV1ToolPhase.call,
},
} satisfies StreamEvent,
context,
execContext,
{ interactive: false, timeout: 1000 }
)

await flushMacrotask()

expect(executeTool).toHaveBeenCalledWith(
'azure_devops_create_work_item',
{ title: 'Ship it', activity: 'Development' },
expect.any(Object)
)
})

it('clears pending continuation state when a run resumes', async () => {
context.awaitingAsyncContinuation = {
checkpointId: 'cp-1',
Expand Down
16 changes: 13 additions & 3 deletions apps/sim/lib/mothership/request/handlers/tool.ts
Original file line number Diff line number Diff line change
Expand Up @@ -20,6 +20,7 @@ import {
MothershipStreamV1ToolOutcome,
type MothershipStreamV1ToolResultPayload,
} from '@/lib/mothership/generated/mothership-stream-v1'
import type { ToolActivity } from '@/lib/mothership/generated/protocol'
import { TraceAttr } from '@/lib/mothership/generated/trace-attributes-v1'
import { TraceSpan } from '@/lib/mothership/generated/trace-spans-v1'
import { withCopilotSpan } from '@/lib/mothership/request/otel'
Expand Down Expand Up @@ -145,6 +146,11 @@ function applyToolDisplay(toolCall: ToolCallState | undefined): void {

const INTEGRATION_GATEWAY_TOOL = 'call_integration_tool'

/** A call's display grouping: carried in its own arguments, or inherited and held beside them. */
function toolCallActivity(toolCall: ToolCallState | undefined): ToolActivity | undefined {
return toolCall?.activity ?? readToolActivity(toolCall?.params)
}

function handleToolArgsDelta(
data: { argumentsDelta: string; toolCallId: string; toolName: string },
context: StreamingContext
Expand Down Expand Up @@ -579,15 +585,19 @@ async function handleCallPhase(
existing.activityDescription ??= activityDescription
}
/** Persist visible labels even when the introducing discovery call is omitted from the transcript. */
const suppliedActivity = readToolActivity(args) ?? readToolActivity(existing?.params)
const argsActivity = readToolActivity(args)
const suppliedActivity = argsActivity ?? toolCallActivity(existing)
const activity = suppliedActivity?.completedTitle
? suppliedActivity
: suppliedActivity &&
Array.from(context.toolCalls.values())
.filter((tool) => tool.agentId === agentId && tool.parentToolCallId === parentToolCallId)
.map((tool) => readToolActivity(tool.params))
.map(toolCallActivity)
.find((candidate) => candidate?.id === suppliedActivity.id && candidate.completedTitle)
const displayArgs = activity ? { ...args, activity } : args
// Only a label the frame's own arguments carry is completed in place. An inherited one stays
// off the arguments, which execute as given (an MCP schema refuses keys it never declared).
const displayArgs = activity && argsActivity ? { ...args, activity } : args
if (existing && activity && !argsActivity) existing.activity = activity
const isSubagent = scope === 'subagent'
const ui = getToolCallUI(data)

Expand Down
8 changes: 7 additions & 1 deletion apps/sim/lib/mothership/request/types.ts
Loading