Add $select filtering and ref-aware schemas for tool chaining

.changeset/adequate-bronze-smelt.md

@@ -0,0 +1,5 @@
+---
+"@inkeep/agents-api": patch
+---
+
+Improve tool chaining with ref-aware schemas, better  error messages, and shared schema description constants

agents-api/src/__tests__/run/agents/PromptConfig.artifactSchema.test.ts

@@ -85,7 +85,7 @@ describe('PromptConfig — artifact type and schema in generated XML', () => {
     const typeSchemaMatch = result.prompt.match(/<type_schema>([\s\S]*?)<\/type_schema>/);
     expect(typeSchemaMatch).not.toBeNull();
     expect(typeSchemaMatch?.[1]).toContain('DISPLAYED to user');
-    expect(typeSchemaMatch?.[1]).toContain('PASSED to tools');
+    expect(typeSchemaMatch?.[1]).toContain('TOOL CHAINING');
   });
 
   test('type_schema preview contains only inPreview fields', () => {
@@ -101,7 +101,7 @@ describe('PromptConfig — artifact type and schema in generated XML', () => {
     const typeSchemaMatch = result.prompt.match(/<type_schema>([\s\S]*?)<\/type_schema>/);
     const typeSchemaContent = typeSchemaMatch?.[1] ?? '';
     const previewIndex = typeSchemaContent.indexOf('DISPLAYED to user');
-    const fullIndex = typeSchemaContent.indexOf('PASSED to tools');
+    const fullIndex = typeSchemaContent.indexOf('TOOL CHAINING');
     const previewSection = typeSchemaContent.slice(previewIndex, fullIndex);
     expect(previewSection).toContain('"title"');
     expect(previewSection).toContain('"summary"');
@@ -121,7 +121,7 @@ describe('PromptConfig — artifact type and schema in generated XML', () => {
     const result = builder.buildSystemPrompt(config);
     const typeSchemaMatch = result.prompt.match(/<type_schema>([\s\S]*?)<\/type_schema>/);
     const typeSchemaContent = typeSchemaMatch?.[1] ?? '';
-    const fullIndex = typeSchemaContent.indexOf('PASSED to tools');
+    const fullIndex = typeSchemaContent.indexOf('TOOL CHAINING');
     const fullSection = typeSchemaContent.slice(fullIndex);
     expect(fullSection).toContain('"title"');
     expect(fullSection).toContain('"summary"');

agents-api/src/__tests__/run/agents/ref-aware-schema.test.ts

@@ -0,0 +1,152 @@
+import { describe, expect, it } from 'vitest';
+import {
+  makeBaseInputSchema,
+  makeRefAwareJsonSchema,
+} from '../../../domains/run/agents/tools/ref-aware-schema';
+import { SENTINEL_KEY } from '../../../domains/run/constants/artifact-syntax';
+
+describe('makeRefAwareJsonSchema', () => {
+  it('should transform string properties to accept tool refs', () => {
+    const schema = {
+      type: 'object',
+      properties: {
+        text: { type: 'string', description: 'The text to measure' },
+      },
+      required: ['text'],
+    };
+
+    const result = makeRefAwareJsonSchema(schema);
+
+    expect(result.required).toEqual(['text']);
+    const textProp = (result.properties as any).text;
+    expect(textProp.anyOf).toBeDefined();
+    expect(textProp.anyOf).toHaveLength(2);
+    expect(textProp.anyOf[0]).toEqual({ type: 'string', description: 'The text to measure' });
+
+    const refOption = textProp.anyOf[1];
+    expect(refOption.anyOf).toHaveLength(2);
+    expect(refOption.anyOf[0].required).toContain(SENTINEL_KEY.ARTIFACT);
+    expect(refOption.anyOf[0].required).toContain(SENTINEL_KEY.TOOL);
+    expect(refOption.anyOf[1].required).toContain(SENTINEL_KEY.TOOL);
+  });
+
+  it('should transform nested object properties', () => {
+    const schema = {
+      type: 'object',
+      properties: {
+        config: {
+          type: 'object',
+          properties: {
+            name: { type: 'string' },
+            count: { type: 'number' },
+          },
+        },
+      },
+    };
+
+    const result = makeRefAwareJsonSchema(schema);
+
+    const configProp = (result.properties as any).config;
+    expect(configProp.anyOf).toBeDefined();
+
+    const configObj = configProp.anyOf[0];
+    expect((configObj.properties as any).name.anyOf).toBeDefined();
+    expect((configObj.properties as any).count.anyOf).toBeDefined();
+  });
+
+  it('should transform array item types', () => {
+    const schema = {
+      type: 'object',
+      properties: {
+        items: {
+          type: 'array',
+          items: { type: 'string' },
+        },
+      },
+    };
+
+    const result = makeRefAwareJsonSchema(schema);
+
+    const itemsProp = (result.properties as any).items;
+    expect(itemsProp.anyOf).toBeDefined();
+    const arraySchema = itemsProp.anyOf[0];
+    expect(arraySchema.items.anyOf).toBeDefined();
+  });
+
+  it('should include $tool, $select, and $artifact in ref schema', () => {
+    const schema = {
+      type: 'object',
+      properties: {
+        text: { type: 'string' },
+      },
+    };
+
+    const result = makeRefAwareJsonSchema(schema);
+    const textProp = (result.properties as any).text;
+    const refSchema = textProp.anyOf[1];
+
+    const toolOnlyRef = refSchema.anyOf[1];
+    expect(toolOnlyRef.properties).toHaveProperty(SENTINEL_KEY.TOOL);
+    expect(toolOnlyRef.properties).toHaveProperty(SENTINEL_KEY.SELECT);
+
+    const artifactRef = refSchema.anyOf[0];
+    expect(artifactRef.properties).toHaveProperty(SENTINEL_KEY.ARTIFACT);
+    expect(artifactRef.properties).toHaveProperty(SENTINEL_KEY.TOOL);
+    expect(artifactRef.properties).toHaveProperty(SENTINEL_KEY.SELECT);
+  });
+
+  it('should not transform the root schema itself', () => {
+    const schema = {
+      type: 'object',
+      properties: {
+        text: { type: 'string' },
+      },
+    };
+
+    const result = makeRefAwareJsonSchema(schema);
+    expect(result.type).toBe('object');
+    expect(result.anyOf).toBeUndefined();
+  });
+
+  it('should handle anyOf/oneOf variants without adding refs', () => {
+    const schema = {
+      type: 'object',
+      properties: {
+        value: {
+          anyOf: [{ type: 'string' }, { type: 'number' }],
+        },
+      },
+    };
+
+    const result = makeRefAwareJsonSchema(schema);
+    const valueProp = (result.properties as any).value;
+    expect(valueProp.anyOf).toBeDefined();
+  });
+
+  it('should handle empty schema gracefully', () => {
+    const schema = { type: 'object' };
+    const result = makeRefAwareJsonSchema(schema);
+    expect(result.type).toBe('object');
+  });
+});
+
+describe('makeBaseInputSchema', () => {
+  it('should return a zod schema from JSON schema', () => {
+    const schema = {
+      type: 'object',
+      properties: {
+        text: { type: 'string' },
+      },
+      required: ['text'],
+    };
+
+    const zodSchema = makeBaseInputSchema(schema);
+    expect(zodSchema.safeParse).toBeDefined();
+
+    const valid = zodSchema.safeParse({ text: 'hello' });
+    expect(valid.success).toBe(true);
+
+    const invalid = zodSchema.safeParse({ text: 123 });
+    expect(invalid.success).toBe(false);
+  });
+});

agents-api/src/__tests__/run/artifacts/resolveArgs-select.test.ts

@@ -0,0 +1,190 @@
+import type { FullExecutionContext } from '@inkeep/agents-core';
+import { describe, expect, it, vi } from 'vitest';
+import {
+  ArtifactParser,
+  ToolChainResolutionError,
+} from '../../../domains/run/artifacts/ArtifactParser';
+import type { ArtifactService } from '../../../domains/run/artifacts/ArtifactService';
+
+function createParser(overrides: {
+  getArtifactFull?: ArtifactService['getArtifactFull'];
+  getToolResultRaw?: ArtifactService['getToolResultRaw'];
+  getToolResultFull?: ArtifactService['getToolResultFull'];
+}) {
+  const mockArtifactService = {
+    getArtifactFull: overrides.getArtifactFull ?? vi.fn().mockResolvedValue(null),
+    getToolResultRaw: overrides.getToolResultRaw ?? vi.fn().mockReturnValue(undefined),
+    getToolResultFull:
+      overrides.getToolResultFull ??
+      overrides.getToolResultRaw ??
+      vi.fn().mockReturnValue(undefined),
+  } as unknown as ArtifactService;
+
+  const mockExecContext = {
+    tenantId: 'test-tenant',
+    projectId: 'test-project',
+  } as FullExecutionContext;
+
+  return new ArtifactParser(mockExecContext, {
+    artifactService: mockArtifactService,
+  });
+}
+
+describe('resolveArgs with $select', () => {
+  const toolData = {
+    items: [
+      { title: 'Alpha', score: 0.9 },
+      { title: 'Beta', score: 0.5 },
+      { title: 'Gamma', score: 0.85 },
+    ],
+    metadata: { total: 3 },
+  };
+
+  describe('$tool + $select', () => {
+    it('should filter ephemeral tool result with $select', async () => {
+      const parser = createParser({
+        getToolResultRaw: vi.fn().mockReturnValue(toolData),
+      });
+
+      const result = await parser.resolveArgs({
+        $tool: 'call_search',
+        $select: 'items[?score > `0.8`]',
+      });
+
+      expect(result).toEqual([
+        { title: 'Alpha', score: 0.9 },
+        { title: 'Gamma', score: 0.85 },
+      ]);
+    });
+
+    it('should extract specific fields with $select', async () => {
+      const parser = createParser({
+        getToolResultRaw: vi.fn().mockReturnValue(toolData),
+      });
+
+      const result = await parser.resolveArgs({
+        $tool: 'call_search',
+        $select: 'items[].title',
+      });
+
+      expect(result).toEqual(['Alpha', 'Beta', 'Gamma']);
+    });
+
+    it('should throw ToolChainResolutionError when $select matches nothing', async () => {
+      const parser = createParser({
+        getToolResultRaw: vi.fn().mockReturnValue(toolData),
+      });
+
+      await expect(
+        parser.resolveArgs({
+          $tool: 'call_search',
+          $select: 'nonexistent_field',
+        })
+      ).rejects.toThrow(ToolChainResolutionError);
+    });
+  });
+
+  describe('$artifact + $tool + $select', () => {
+    it('should filter artifact data with $select', async () => {
+      const parser = createParser({
+        getArtifactFull: vi.fn().mockResolvedValue({ data: toolData }),
+      });
+
+      const result = await parser.resolveArgs({
+        $artifact: 'art_123',
+        $tool: 'call_search',
+        $select: 'metadata.total',
+      });
+
+      expect(result).toBe(3);
+    });
+  });
+
+  describe('regression: without $select', () => {
+    it('should return full tool result when no $select', async () => {
+      const parser = createParser({
+        getToolResultRaw: vi.fn().mockReturnValue(toolData),
+      });
+
+      const result = await parser.resolveArgs({
+        $tool: 'call_search',
+      });
+
+      expect(result).toEqual(toolData);
+    });
+
+    it('should return full artifact data when no $select', async () => {
+      const parser = createParser({
+        getArtifactFull: vi.fn().mockResolvedValue({ data: toolData }),
+      });
+
+      const result = await parser.resolveArgs({
+        $artifact: 'art_123',
+        $tool: 'call_search',
+      });
+
+      expect(result).toEqual(toolData);
+    });
+  });
+
+  describe('error handling', () => {
+    it('should throw ToolChainResolutionError for bad JMESPath', async () => {
+      const parser = createParser({
+        getToolResultRaw: vi.fn().mockReturnValue(toolData),
+      });
+
+      await expect(
+        parser.resolveArgs({
+          $tool: 'call_search',
+          $select: '[invalid!!!',
+        })
+      ).rejects.toThrow(ToolChainResolutionError);
+    });
+
+    it('should throw ToolChainResolutionError for dangerous expression', async () => {
+      const parser = createParser({
+        getToolResultRaw: vi.fn().mockReturnValue(toolData),
+      });
+
+      await expect(
+        parser.resolveArgs({
+          $tool: 'call_search',
+          $select: '__proto__',
+        })
+      ).rejects.toThrow(ToolChainResolutionError);
+    });
+  });
+
+  describe('nested refs with $select', () => {
+    it('should resolve nested $select at different levels', async () => {
+      const parser = createParser({
+        getToolResultRaw: vi.fn().mockReturnValue(toolData),
+      });
+
+      const result = await parser.resolveArgs({
+        first: { $tool: 'call_search', $select: 'items[0].title' },
+        second: { $tool: 'call_search', $select: 'metadata.total' },
+      });
+
+      expect(result).toEqual({
+        first: 'Alpha',
+        second: 3,
+      });
+    });
+  });
+
+  describe('result. prefix stripping', () => {
+    it('should auto-strip result. prefix from $select expressions', async () => {
+      const parser = createParser({
+        getToolResultRaw: vi.fn().mockReturnValue(toolData),
+      });
+
+      const result = await parser.resolveArgs({
+        $tool: 'call_search',
+        $select: 'result.metadata.total',
+      });
+
+      expect(result).toBe(3);
+    });
+  });
+});