Skip to content

Commit 7137481

Browse files
fix(tables): pass enriched query schema to agents
1 parent 117fe31 commit 7137481

18 files changed

Lines changed: 387 additions & 57 deletions

apps/sim/executor/handlers/agent/agent-handler.test.ts

Lines changed: 19 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -29,6 +29,7 @@ import { SIM_AUTO_MODEL_ID } from '@/providers/models'
2929
import { getProviderFromModel, transformBlockTool } from '@/providers/utils'
3030
import type { SerializedBlock, SerializedWorkflow } from '@/serializer/types'
3131
import { executeTool } from '@/tools'
32+
import { ToolSchemaEnrichmentError } from '@/tools/params'
3233

3334
process.env.NEXT_PUBLIC_APP_URL = 'http://localhost:3000'
3435

@@ -289,6 +290,24 @@ describe('AgentBlockHandler', () => {
289290
expect(result).toEqual(expectedOutput)
290291
})
291292

293+
it('fails fast when a configured tool schema cannot be enriched', async () => {
294+
const error = new ToolSchemaEnrichmentError(
295+
'table_query_rows',
296+
new Error('table metadata unavailable')
297+
)
298+
mockTransformBlockTool.mockRejectedValueOnce(error)
299+
300+
await expect(
301+
handler.execute(mockContext, mockBlock, {
302+
model: 'gpt-4o',
303+
userPrompt: 'Query the table',
304+
apiKey: 'test-api-key',
305+
tools: [{ type: 'table', operation: 'query_rows', usageControl: 'auto' }],
306+
})
307+
).rejects.toBe(error)
308+
expect(mockExecuteProviderRequest).not.toHaveBeenCalled()
309+
})
310+
292311
it('reports a sim-auto run under the sim-auto identity, not the model that served it', async () => {
293312
mockExecuteProviderRequest.mockResolvedValue({
294313
content: 'Mocked response content',

apps/sim/executor/handlers/agent/agent-handler.ts

Lines changed: 35 additions & 1 deletion
Original file line numberDiff line numberDiff line change
@@ -6,6 +6,7 @@ import { sleep } from '@sim/utils/helpers'
66
import { isPlainRecord } from '@sim/utils/object'
77
import { truncate } from '@sim/utils/string'
88
import { and, eq, inArray, isNull } from 'drizzle-orm'
9+
import { isDev } from '@/lib/core/config/env-flags'
910
import { normalizeStringRecord, normalizeWorkflowVariables } from '@/lib/core/utils/records'
1011
import { createMcpToolId } from '@/lib/mcp/utils'
1112
import {
@@ -67,7 +68,7 @@ import {
6768
import { isAutoModel, SIM_AUTO_MODEL_ID } from '@/providers/models'
6869
import { getProviderFromModel, transformBlockTool } from '@/providers/utils'
6970
import type { SerializedBlock } from '@/serializer/types'
70-
import { filterSchemaForLLM, type ToolSchema } from '@/tools/params'
71+
import { filterSchemaForLLM, type ToolSchema, ToolSchemaEnrichmentError } from '@/tools/params'
7172
import { getTool } from '@/tools/utils'
7273
import { getToolAsync } from '@/tools/utils.server'
7374

@@ -247,6 +248,32 @@ export class AgentBlockHandler implements BlockHandler {
247248
streaming: streamingConfig.shouldUseStreaming ?? false,
248249
})
249250

251+
if (isDev) {
252+
const tableQueryRowsV1Tools = formattedTools.filter((tool) => {
253+
const toolId = tool?.id
254+
if (typeof toolId !== 'string') return false
255+
256+
const isTableQueryRows =
257+
toolId === 'table_query_rows' || toolId.startsWith('table_query_rows_')
258+
const isV2 = toolId === 'table_query_rows_v2' || toolId.startsWith('table_query_rows_v2_')
259+
return isTableQueryRows && !isV2
260+
})
261+
262+
if (tableQueryRowsV1Tools.length > 0) {
263+
logger.info('Passing table_query_rows v1 tool schema to agent provider', {
264+
blockId: block.id,
265+
executionId: ctx.executionId,
266+
providerId,
267+
model,
268+
tools: tableQueryRowsV1Tools.map((tool) => ({
269+
id: tool.id,
270+
description: tool.description,
271+
parameters: tool.parameters,
272+
})),
273+
})
274+
}
275+
}
276+
250277
const result = await this.executeProviderRequest(ctx, providerRequest, block, responseFormat)
251278

252279
if (autoRouting && autoRouting.billableRoutingCost > 0) {
@@ -526,6 +553,7 @@ export class AgentBlockHandler implements BlockHandler {
526553
}
527554
return this.transformBlockTool(ctx, tool, canonicalModes, toolIndex)
528555
} catch (error) {
556+
if (error instanceof ToolSchemaEnrichmentError) throw error
529557
logger.error(
530558
'[AgentHandler] Error creating tool',
531559
projectAgentDiagnosticMetadata(
@@ -952,6 +980,12 @@ export class AgentBlockHandler implements BlockHandler {
952980
}),
953981
getTool,
954982
canonicalModes,
983+
enrichmentContext: {
984+
workflowId: ctx.workflowId,
985+
workspaceId: ctx.workspaceId,
986+
executionId: ctx.executionId,
987+
userId: ctx.userId,
988+
},
955989
toolIndex,
956990
resolveCustomBlockBinding: (blockType: string) =>
957991
resolveCustomBlockToolBinding(blockType, ctx.workspaceId),

apps/sim/executor/handlers/pi/sim-tools.test.ts

Lines changed: 15 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -16,6 +16,7 @@ vi.mock('@/tools/utils.server', () => ({ getToolAsync: vi.fn() }))
1616
import { buildSimToolSpecs } from '@/executor/handlers/pi/sim-tools'
1717
import type { ExecutionContext } from '@/executor/types'
1818
import { ResolvedSecretTraceRegistry } from '@/executor/utils/resolved-secret-trace-registry'
19+
import { ToolSchemaEnrichmentError } from '@/tools/params'
1920

2021
function executionContext(registry: ResolvedSecretTraceRegistry | undefined): ExecutionContext {
2122
return {
@@ -76,6 +77,20 @@ describe('buildSimToolSpecs', () => {
7677
expect(mockTransformBlockTool).not.toHaveBeenCalled()
7778
})
7879

80+
it('fails fast when a tool schema cannot be enriched', async () => {
81+
const error = new ToolSchemaEnrichmentError(
82+
'table_query_rows',
83+
new Error('table metadata unavailable')
84+
)
85+
mockTransformBlockTool.mockRejectedValueOnce(error)
86+
87+
await expect(
88+
buildSimToolSpecs(completeExecutionContext(), [
89+
{ type: 'table', operation: 'query_rows', usageControl: 'auto' },
90+
])
91+
).rejects.toBe(error)
92+
})
93+
7994
it('forwards a trusted _context that an LLM-supplied _context cannot override', async () => {
8095
mockTransformBlockTool.mockResolvedValue({
8196
id: 'exa_search',

apps/sim/executor/handlers/pi/sim-tools.ts

Lines changed: 8 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -22,6 +22,7 @@ import type { ResolvedSecretTraceRegistry } from '@/executor/utils/resolved-secr
2222
import { transformBlockTool } from '@/providers/utils'
2323
import { executeTool } from '@/tools'
2424
import { mergeToolParameters } from '@/tools/merge-params'
25+
import { ToolSchemaEnrichmentError } from '@/tools/params'
2526
import type { ToolResponse } from '@/tools/types'
2627
import { getTool } from '@/tools/utils'
2728
import { getToolAsync } from '@/tools/utils.server'
@@ -97,6 +98,12 @@ export async function buildSimToolSpecs(
9798
getAllBlocks,
9899
getTool,
99100
getToolAsync,
101+
enrichmentContext: {
102+
workflowId: ctx.workflowId,
103+
workspaceId: ctx.workspaceId,
104+
executionId: ctx.executionId,
105+
userId: ctx.userId,
106+
},
100107
resolveCustomBlockBinding: (blockType: string) =>
101108
resolveCustomBlockToolBinding(blockType, ctx.workspaceId),
102109
})
@@ -171,6 +178,7 @@ export async function buildSimToolSpecs(
171178
},
172179
})
173180
} catch (error) {
181+
if (error instanceof ToolSchemaEnrichmentError) throw error
174182
logger.warn('Failed to adapt Sim tool for Pi', {
175183
type: tool.type,
176184
error: getErrorMessage(error),

apps/sim/providers/utils.test.ts

Lines changed: 72 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -1622,6 +1622,78 @@ describe('transformBlockTool multi-instance unique IDs', () => {
16221622
expect(result?.id).toBe('table_query_rows_tbl_abc')
16231623
})
16241624

1625+
it('resolves the canonical table id before enriching the LLM tool schema', async () => {
1626+
const enrichTool = vi.fn(
1627+
async (
1628+
tableId: string,
1629+
schema: {
1630+
type: 'object'
1631+
properties: Record<string, unknown>
1632+
required: string[]
1633+
}
1634+
) => ({
1635+
description: `Query rows from ${tableId}`,
1636+
parameters: {
1637+
...schema,
1638+
properties: {
1639+
...schema.properties,
1640+
customer_name: { type: 'string' },
1641+
},
1642+
},
1643+
})
1644+
)
1645+
const result = await transformBlockTool(
1646+
{
1647+
type: 'table',
1648+
operation: 'query_rows',
1649+
params: { tableSelector: 'tbl_abc' },
1650+
},
1651+
{
1652+
selectedOperation: 'query_rows',
1653+
getAllBlocks,
1654+
enrichmentContext: {
1655+
workspaceId: 'workspace-1',
1656+
userId: 'user-1',
1657+
},
1658+
getTool: (id: string) => ({
1659+
id,
1660+
name: 'Query Rows',
1661+
description: 'Query table rows',
1662+
params: {
1663+
tableId: { type: 'string', required: true, visibility: 'user-only' },
1664+
filter: { type: 'object', visibility: 'user-or-llm' },
1665+
},
1666+
toolEnrichment: {
1667+
dependsOn: 'tableId',
1668+
enrichTool,
1669+
},
1670+
}),
1671+
}
1672+
)
1673+
1674+
expect(enrichTool).toHaveBeenCalledWith(
1675+
'tbl_abc',
1676+
expect.objectContaining({
1677+
properties: expect.objectContaining({ filter: expect.any(Object) }),
1678+
}),
1679+
'Query table rows',
1680+
{
1681+
workspaceId: 'workspace-1',
1682+
userId: 'user-1',
1683+
}
1684+
)
1685+
expect(result).toMatchObject({
1686+
id: 'table_query_rows_tbl_abc',
1687+
description: 'Query rows from tbl_abc',
1688+
params: { tableSelector: 'tbl_abc' },
1689+
parameters: {
1690+
properties: {
1691+
customer_name: { type: 'string' },
1692+
},
1693+
},
1694+
})
1695+
})
1696+
16251697
it('appends the table id resolved from the advanced manual input', async () => {
16261698
const result = await transformTable(
16271699
{ manualTableId: 'tbl_xyz' },

apps/sim/providers/utils.ts

Lines changed: 28 additions & 9 deletions
Original file line numberDiff line numberDiff line change
@@ -5,7 +5,7 @@ import type OpenAI from 'openai'
55
import type { BillingAttributionSnapshot } from '@/lib/billing/core/billing-attribution'
66
import { formatCreditCost } from '@/lib/billing/credits/conversion'
77
import { env } from '@/lib/core/config/env'
8-
import { getBlacklistedProvidersFromEnv, isHosted } from '@/lib/core/config/env-flags'
8+
import { getBlacklistedProvidersFromEnv, isDev, isHosted } from '@/lib/core/config/env-flags'
99
import {
1010
normalizeRecord,
1111
normalizeStringRecord,
@@ -53,6 +53,7 @@ import {
5353
import type { ProviderId, ProviderToolConfig } from '@/providers/types'
5454
import { useProvidersStore } from '@/stores/providers/store'
5555
import { mergeToolParameters } from '@/tools/merge-params'
56+
import type { WorkflowToolExecutionContext } from '@/tools/types'
5657

5758
const logger = createLogger('ProviderUtils')
5859

@@ -629,6 +630,7 @@ export async function transformBlockTool(
629630
getTool: (toolId: string) => any
630631
getToolAsync?: (toolId: string) => Promise<any>
631632
canonicalModes?: Record<string, 'basic' | 'advanced'>
633+
enrichmentContext?: WorkflowToolExecutionContext
632634
/**
633635
* Server-only resolver for a custom (deploy-as-block) tool's binding (bound
634636
* workflow + input schema), org-scoped to the consumer. Injected as a dependency
@@ -646,8 +648,15 @@ export async function transformBlockTool(
646648
toolIndex?: number
647649
}
648650
): Promise<ProviderToolConfig | null> {
649-
const { selectedOperation, getAllBlocks, getTool, getToolAsync, canonicalModes, toolIndex } =
650-
options
651+
const {
652+
selectedOperation,
653+
getAllBlocks,
654+
getTool,
655+
getToolAsync,
656+
canonicalModes,
657+
enrichmentContext,
658+
toolIndex,
659+
} = options
651660
const scopedCanonicalModes = scopeCanonicalModesForTool(canonicalModes, toolIndex, block.type)
652661

653662
const blockDef = getAllBlocks().find((b: any) => b.type === block.type)
@@ -755,12 +764,6 @@ export async function transformBlockTool(
755764

756765
const userProvidedParams = block.params || {}
757766

758-
const {
759-
schema: llmSchema,
760-
enrichedDescription,
761-
modelBlockedParams,
762-
} = await createLLMToolSchema(toolConfig, userProvidedParams)
763-
764767
const canonicalGroups: CanonicalGroup[] = blockDef?.subBlocks
765768
? Object.values(buildCanonicalIndex(blockDef.subBlocks).groupsById).filter(isCanonicalPair)
766769
: []
@@ -771,6 +774,12 @@ export async function transformBlockTool(
771774
scopedCanonicalModes
772775
)
773776

777+
const {
778+
schema: llmSchema,
779+
enrichedDescription,
780+
modelBlockedParams,
781+
} = await createLLMToolSchema(toolConfig, resolvedResourceParams, enrichmentContext)
782+
774783
let uniqueToolId = toolConfig.id
775784
let toolName = toolConfig.name
776785
let toolDescription = enrichedDescription || toolConfig.description
@@ -857,6 +866,16 @@ export async function transformBlockTool(
857866
}
858867
: undefined
859868

869+
if (isDev && toolConfig.id === 'table_query_rows') {
870+
logger.info('Prepared table_query_rows v1 tool schema', {
871+
toolId: uniqueToolId,
872+
tableId: resolvedResourceParams.tableId,
873+
enrichmentApplied: Boolean(enrichedDescription),
874+
description: toolDescription,
875+
parameters: llmSchema,
876+
})
877+
}
878+
860879
return {
861880
id: uniqueToolId,
862881
name: toolName,

apps/sim/tools/params.test.ts

Lines changed: 22 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -12,6 +12,7 @@ import {
1212
isPasswordParameter,
1313
type ToolParameterConfig,
1414
type ToolSchema,
15+
ToolSchemaEnrichmentError,
1516
type ValidationResult,
1617
validateToolParameters,
1718
} from '@/tools/params'
@@ -130,6 +131,27 @@ describe('Tool Parameters Utils', () => {
130131
expect(schema.required).not.toContain('apiKey') // user-only, never required for LLM
131132
expect(schema.required).toContain('message') // user-or-llm + required: true
132133
})
134+
135+
it('wraps tool enrichment failures so execution boundaries can fail fast', async () => {
136+
const cause = new Error('table metadata unavailable')
137+
const toolConfig = {
138+
...mockToolConfig,
139+
toolEnrichment: {
140+
dependsOn: 'tableId',
141+
enrichTool: vi.fn().mockRejectedValue(cause),
142+
},
143+
}
144+
145+
const error = await createLLMToolSchema(toolConfig, { tableId: 'tbl_123' }).catch(
146+
(caught) => caught
147+
)
148+
149+
expect(error).toBeInstanceOf(ToolSchemaEnrichmentError)
150+
expect(error).toMatchObject({
151+
message: 'Failed to enrich schema for tool "test_tool"',
152+
cause,
153+
})
154+
})
133155
})
134156

135157
describe('createUserToolSchema', () => {

0 commit comments

Comments
 (0)