Skip to content

Commit bfbc4f6

Browse files
authored
feat(toolbox-core): Add support for optional parameters (#66)
* feat(toolbox-core): Add support for optional parameters * chore: Delint * chore: Add unit test cases for protocol file * chore: Delint * chore: Add unit tests to test null value in request payload * chore: Add e2e tests for optional params * chore: Fix e2e tests type assertion * chore: Fix some e2e tests * chore: Fix e2e tests
1 parent 9e60069 commit bfbc4f6

File tree

5 files changed

+318
-8
lines changed

5 files changed

+318
-8
lines changed

packages/toolbox-core/src/toolbox_core/protocol.ts

Lines changed: 21 additions & 5 deletions
Original file line numberDiff line numberDiff line change
@@ -20,6 +20,7 @@ interface BaseParameter {
2020
name: string;
2121
description: string;
2222
authSources?: string[];
23+
required?: boolean;
2324
}
2425

2526
interface StringParameter extends BaseParameter {
@@ -56,6 +57,7 @@ const ZodBaseParameter = z.object({
5657
name: z.string().min(1, 'Parameter name cannot be empty'),
5758
description: z.string(),
5859
authSources: z.array(z.string()).optional(),
60+
required: z.boolean().optional(),
5961
});
6062

6163
export const ZodParameterSchema = z.lazy(() =>
@@ -101,24 +103,38 @@ export type ZodManifest = z.infer<typeof ZodManifestSchema>;
101103
* @returns A ZodTypeAny representing the schema for this parameter.
102104
*/
103105
function buildZodShapeFromParam(param: ParameterSchema): ZodTypeAny {
106+
let schema: ZodTypeAny;
104107
switch (param.type) {
105108
case 'string':
106-
return z.string();
109+
schema = z.string();
110+
break;
107111
case 'integer':
108-
return z.number().int();
112+
schema = z.number().int();
113+
break;
109114
case 'float':
110-
return z.number();
115+
schema = z.number();
116+
break;
111117
case 'boolean':
112-
return z.boolean();
118+
schema = z.boolean();
119+
break;
113120
case 'array':
114121
// Recursively build the schema for array items
115-
return z.array(buildZodShapeFromParam(param.items));
122+
// Array items inherit the 'required' status of the parent array.
123+
param.items.required = param.required;
124+
schema = z.array(buildZodShapeFromParam(param.items));
125+
break;
116126
default: {
117127
// This ensures exhaustiveness at compile time if ParameterSchema is a discriminated union
118128
const _exhaustiveCheck: never = param;
119129
throw new Error(`Unknown parameter type: ${_exhaustiveCheck['type']}`);
120130
}
121131
}
132+
133+
if (param.required === false) {
134+
return schema.nullish();
135+
}
136+
137+
return schema;
122138
}
123139

124140
/**

packages/toolbox-core/src/toolbox_core/tool.ts

Lines changed: 18 additions & 3 deletions
Original file line numberDiff line numberDiff line change
@@ -138,6 +138,17 @@ function ToolboxTool(
138138

139139
const payload = {...validatedUserArgs, ...resolvedBoundParams};
140140

141+
// Filter out null values from the payload
142+
const filteredPayload = Object.entries(payload).reduce(
143+
(acc, [key, value]) => {
144+
if (value !== null && value !== undefined) {
145+
acc[key] = value;
146+
}
147+
return acc;
148+
},
149+
{} as Record<string, unknown>
150+
);
151+
141152
const headers: Record<string, string> = {};
142153
for (const [headerName, headerValue] of Object.entries(clientHeaders)) {
143154
const resolvedHeaderValue = await resolveValue(headerValue);
@@ -159,9 +170,13 @@ function ToolboxTool(
159170
}
160171

161172
try {
162-
const response: AxiosResponse = await session.post(toolUrl, payload, {
163-
headers,
164-
});
173+
const response: AxiosResponse = await session.post(
174+
toolUrl,
175+
filteredPayload,
176+
{
177+
headers,
178+
}
179+
);
165180
return typeof response.data === 'string'
166181
? response.data
167182
: JSON.stringify(response.data);

packages/toolbox-core/test/e2e/test.e2e.ts

Lines changed: 177 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -18,6 +18,7 @@ import {ToolboxTool} from '../../src/toolbox_core/tool';
1818
import {AxiosError} from 'axios';
1919
import {CustomGlobal} from './types';
2020
import {authTokenGetter} from './utils';
21+
import {ZodOptional, ZodNullable, ZodTypeAny} from 'zod';
2122

2223
describe('ToolboxClient E2E Tests', () => {
2324
let commonToolboxClient: ToolboxClient;
@@ -316,4 +317,180 @@ describe('ToolboxClient E2E Tests', () => {
316317
}
317318
});
318319
});
320+
321+
describe('Optional Params E2E Tests', () => {
322+
let searchRowsTool: ReturnType<typeof ToolboxTool>;
323+
324+
beforeAll(async () => {
325+
searchRowsTool = await commonToolboxClient.loadTool('search-rows');
326+
});
327+
328+
it('should correctly identify required and optional parameters in the schema', () => {
329+
const paramSchema = searchRowsTool.getParamSchema();
330+
const {shape} = paramSchema;
331+
332+
// Required param 'email'
333+
expect(shape.email.isOptional()).toBe(false);
334+
expect(shape.email.isNullable()).toBe(false);
335+
expect(shape.email._def.typeName).toBe('ZodString');
336+
337+
// Optional param 'data'
338+
expect(shape.data.isOptional()).toBe(true);
339+
expect(shape.data.isNullable()).toBe(true);
340+
expect(
341+
(shape.data as ZodOptional<ZodNullable<ZodTypeAny>>).unwrap().unwrap()
342+
._def.typeName
343+
).toBe('ZodString');
344+
345+
// Optional param 'id'
346+
expect(shape.id.isOptional()).toBe(true);
347+
expect(shape.id.isNullable()).toBe(true);
348+
expect(
349+
(shape.id as ZodOptional<ZodNullable<ZodTypeAny>>).unwrap().unwrap()
350+
._def.typeName
351+
).toBe('ZodNumber');
352+
});
353+
354+
it('should run tool with optional params omitted', async () => {
355+
const response = await searchRowsTool({email: '[email protected]'});
356+
expect(typeof response).toBe('string');
357+
expect(response).toContain('\\"email\\":\\"[email protected]\\"');
358+
expect(response).not.toContain('row1');
359+
expect(response).toContain('row2');
360+
expect(response).not.toContain('row3');
361+
expect(response).not.toContain('row4');
362+
expect(response).not.toContain('row5');
363+
expect(response).not.toContain('row6');
364+
});
365+
366+
it('should run tool with optional data provided', async () => {
367+
const response = await searchRowsTool({
368+
369+
data: 'row3',
370+
});
371+
expect(typeof response).toBe('string');
372+
expect(response).toContain('\\"email\\":\\"[email protected]\\"');
373+
expect(response).not.toContain('row1');
374+
expect(response).not.toContain('row2');
375+
expect(response).toContain('row3');
376+
expect(response).not.toContain('row4');
377+
expect(response).not.toContain('row5');
378+
expect(response).not.toContain('row6');
379+
});
380+
381+
it('should run tool with optional data as null', async () => {
382+
const response = await searchRowsTool({
383+
384+
data: null,
385+
});
386+
expect(typeof response).toBe('string');
387+
expect(response).toContain('\\"email\\":\\"[email protected]\\"');
388+
expect(response).not.toContain('row1');
389+
expect(response).toContain('row2');
390+
expect(response).not.toContain('row3');
391+
expect(response).not.toContain('row4');
392+
expect(response).not.toContain('row5');
393+
expect(response).not.toContain('row6');
394+
});
395+
396+
it('should run tool with optional id provided', async () => {
397+
const response = await searchRowsTool({
398+
399+
id: 1,
400+
});
401+
expect(typeof response).toBe('string');
402+
expect(response).toBe('{"result":"null"}');
403+
});
404+
405+
it('should run tool with optional id as null', async () => {
406+
const response = await searchRowsTool({
407+
408+
id: null,
409+
});
410+
expect(typeof response).toBe('string');
411+
expect(response).toContain('\\"email\\":\\"[email protected]\\"');
412+
expect(response).not.toContain('row1');
413+
expect(response).toContain('row2');
414+
expect(response).not.toContain('row3');
415+
expect(response).not.toContain('row4');
416+
expect(response).not.toContain('row5');
417+
expect(response).not.toContain('row6');
418+
});
419+
420+
it('should fail when a required param is missing', async () => {
421+
await expect(searchRowsTool({id: 5, data: 'row5'})).rejects.toThrow(
422+
/Argument validation failed for tool "search-rows":\s*- email: Required/
423+
);
424+
});
425+
426+
it('should fail when a required param is null', async () => {
427+
await expect(
428+
searchRowsTool({email: null, id: 5, data: 'row5'})
429+
).rejects.toThrow(
430+
/Argument validation failed for tool "search-rows":\s*- email: Expected string, received null/
431+
);
432+
});
433+
434+
it('should run tool with all default params', async () => {
435+
const response = await searchRowsTool({
436+
437+
id: 0,
438+
data: 'row2',
439+
});
440+
expect(typeof response).toBe('string');
441+
expect(response).toContain('\\"email\\":\\"[email protected]\\"');
442+
expect(response).not.toContain('row1');
443+
expect(response).toContain('row2');
444+
expect(response).not.toContain('row3');
445+
expect(response).not.toContain('row4');
446+
expect(response).not.toContain('row5');
447+
expect(response).not.toContain('row6');
448+
});
449+
450+
it('should run tool with all valid params', async () => {
451+
const response = await searchRowsTool({
452+
453+
id: 3,
454+
data: 'row3',
455+
});
456+
expect(typeof response).toBe('string');
457+
expect(response).toContain('\\"email\\":\\"[email protected]\\"');
458+
expect(response).not.toContain('row1');
459+
expect(response).not.toContain('row2');
460+
expect(response).toContain('row3');
461+
expect(response).not.toContain('row4');
462+
expect(response).not.toContain('row5');
463+
expect(response).not.toContain('row6');
464+
});
465+
466+
it('should return null when called with a different email', async () => {
467+
const response = await searchRowsTool({
468+
469+
id: 3,
470+
data: 'row3',
471+
});
472+
expect(typeof response).toBe('string');
473+
expect(response).toBe('{"result":"null"}');
474+
});
475+
476+
it('should return null when called with different data', async () => {
477+
const response = await searchRowsTool({
478+
479+
id: 3,
480+
data: 'row4',
481+
});
482+
expect(typeof response).toBe('string');
483+
expect(response).toBe('{"result":"null"}');
484+
});
485+
486+
it('should return null when called with a different id', async () => {
487+
const response = await searchRowsTool({
488+
489+
id: 4,
490+
data: 'row3',
491+
});
492+
expect(typeof response).toBe('string');
493+
expect(response).toBe('{"result":"null"}');
494+
});
495+
});
319496
});

packages/toolbox-core/test/test.protocol.ts

Lines changed: 70 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -119,6 +119,42 @@ describe('ZodParameterSchema', () => {
119119
},
120120
},
121121
},
122+
{
123+
description: 'string parameter with required set to false',
124+
data: {
125+
name: 'optionalString',
126+
description: 'An optional string',
127+
type: 'string',
128+
required: false,
129+
},
130+
},
131+
{
132+
description: 'string parameter with required set to true',
133+
data: {
134+
name: 'requiredString',
135+
description: 'A required string',
136+
type: 'string',
137+
required: true,
138+
},
139+
},
140+
{
141+
description: 'integer parameter with required set to false',
142+
data: {
143+
name: 'optionalInt',
144+
description: 'An optional integer',
145+
type: 'integer',
146+
required: false,
147+
},
148+
},
149+
{
150+
description: 'integer parameter with required set to true',
151+
data: {
152+
name: 'requiredInt',
153+
description: 'A required integer',
154+
type: 'integer',
155+
required: true,
156+
},
157+
},
122158
];
123159

124160
test.each(validParameterTestCases)(
@@ -388,4 +424,38 @@ describe('createZodObjectSchemaFromParameters', () => {
388424
'Unknown parameter type: someUnrecognizedType'
389425
);
390426
});
427+
428+
describe('optional parameters', () => {
429+
const params: ParameterSchema[] = [
430+
{name: 'requiredParam', description: 'required', type: 'string' as const},
431+
{
432+
name: 'optionalParam',
433+
description: 'optional',
434+
type: 'string' as const,
435+
required: false,
436+
},
437+
];
438+
const schema = createZodSchemaFromParams(params);
439+
440+
it('should fail if a required parameter is missing', () => {
441+
expectParseFailure(schema, {optionalParam: 'value'}, errors => {
442+
expect(errors).toContain('requiredParam: Required');
443+
});
444+
});
445+
446+
it('should succeed if an optional parameter is missing', () => {
447+
expectParseSuccess(schema, {requiredParam: 'value'});
448+
});
449+
450+
it('should succeed if an optional parameter is null', () => {
451+
expectParseSuccess(schema, {requiredParam: 'value', optionalParam: null});
452+
});
453+
454+
it('should succeed if an optional parameter is undefined', () => {
455+
expectParseSuccess(schema, {
456+
requiredParam: 'value',
457+
optionalParam: undefined,
458+
});
459+
});
460+
});
391461
});

0 commit comments

Comments
 (0)