Skip to content

Commit 40d0d68

Browse files
bhoustonbhouston
committed
Merge branch 'async-subagents' into interactive-tool-agent
2 parents af661d1 + 6457311 commit 40d0d68

File tree

5 files changed

+460
-2
lines changed

5 files changed

+460
-2
lines changed

‎src/tools/getTools.ts‎

Lines changed: 5 additions & 2 deletions
Original file line numberDiff line numberDiff line change
@@ -1,4 +1,3 @@
1-
import { subAgentTool } from "../tools/interaction/subAgent.js";
21
import { readFileTool } from "../tools/io/readFile.js";
32
import { userPromptTool } from "../tools/interaction/userPrompt.js";
43
import { sequenceCompleteTool } from "../tools/system/sequenceComplete.js";
@@ -7,10 +6,14 @@ import { Tool } from "../core/types.js";
76
import { updateFileTool } from "./io/updateFile.js";
87
import { shellStartTool } from "./system/shellStart.js";
98
import { shellMessageTool } from "./system/shellMessage.js";
9+
import { subAgentStartTool } from "./interaction/subAgentStart.js";
10+
import { subAgentMessageTool } from "./interaction/subAgentMessage.js";
1011

1112
export async function getTools(): Promise<Tool[]> {
1213
return [
13-
subAgentTool,
14+
//subAgentTool, - remove for now.
15+
subAgentStartTool,
16+
subAgentMessageTool,
1417
readFileTool,
1518
updateFileTool,
1619
//shellExecuteTool, - remove for now.

‎src/tools/interaction/subAgentMessage.ts‎

Lines changed: 125 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -0,0 +1,125 @@
1+
import Anthropic from "@anthropic-ai/sdk";
2+
import { Tool } from "../../core/types.js";
3+
import { z } from "zod";
4+
import { zodToJsonSchema } from "zod-to-json-schema";
5+
import { subAgentStates } from "./subAgentState.js";
6+
7+
const parameterSchema = z.object({
8+
instanceId: z
9+
.string()
10+
.describe("The instance ID of the sub-agent to interact with"),
11+
message: z
12+
.string()
13+
.optional()
14+
.describe("The message to send to the sub-agent (required unless aborting)"),
15+
description: z
16+
.string()
17+
.max(80)
18+
.describe("Brief description of the interaction purpose (max 80 chars)"),
19+
abort: z
20+
.boolean()
21+
.optional()
22+
.describe("Whether to abort the sub-agent instead of sending a message"),
23+
});
24+
25+
const returnSchema = z
26+
.object({
27+
response: z.string().describe("Response from the sub-agent or abort confirmation"),
28+
})
29+
.describe("Result containing the sub-agent's response or abort status");
30+
31+
type Parameters = z.infer<typeof parameterSchema>;
32+
type ReturnType = z.infer<typeof returnSchema>;
33+
34+
export const subAgentMessageTool: Tool<Parameters, ReturnType> = {
35+
name: "subAgentMessage",
36+
description: "Sends a message to or aborts an existing sub-agent",
37+
parameters: zodToJsonSchema(parameterSchema),
38+
returns: zodToJsonSchema(returnSchema),
39+
40+
execute: async (
41+
params,
42+
{ logger },
43+
): Promise<ReturnType> => {
44+
if (!process.env.ANTHROPIC_API_KEY) {
45+
throw new Error("ANTHROPIC_API_KEY environment variable is not set");
46+
}
47+
48+
const state = subAgentStates.get(params.instanceId);
49+
if (!state) {
50+
throw new Error("Sub-agent not found");
51+
}
52+
53+
if (state.aborted) {
54+
throw new Error("Sub-agent has been aborted");
55+
}
56+
57+
if (params.abort) {
58+
logger.verbose(`Aborting sub-agent ${params.instanceId}`);
59+
state.aborted = true;
60+
return {
61+
response: "Sub-agent aborted",
62+
};
63+
}
64+
65+
if (!params.message) {
66+
throw new Error("Message is required when not aborting");
67+
}
68+
69+
const anthropic = new Anthropic({
70+
apiKey: process.env.ANTHROPIC_API_KEY,
71+
});
72+
73+
// Create messages array with conversation history
74+
const messages = [
75+
{
76+
role: "user" as const,
77+
content: state.prompt,
78+
},
79+
...state.messages,
80+
{
81+
role: "user" as const,
82+
content: params.message,
83+
},
84+
];
85+
86+
// Get response
87+
logger.verbose(`Sending message to Anthropic API`);
88+
const response = await anthropic.messages.create({
89+
model: "claude-3-opus-20240229",
90+
max_tokens: 4096,
91+
messages,
92+
});
93+
94+
// Check if response content exists and has the expected structure
95+
if (!response.content?.[0]?.text) {
96+
throw new Error("Invalid response from Anthropic API");
97+
}
98+
99+
const responseText = response.content[0].text;
100+
logger.verbose(`Received response from sub-agent`);
101+
102+
// Store the interaction
103+
state.messages.push(
104+
{ role: "user", content: params.message },
105+
{ role: "assistant", content: responseText }
106+
);
107+
108+
return {
109+
response: responseText,
110+
};
111+
},
112+
113+
logParameters: (input, { logger }) => {
114+
if (input.abort) {
115+
logger.info(`Aborting sub-agent ${input.instanceId}: ${input.description}`);
116+
} else {
117+
logger.info(`Sending message to sub-agent ${input.instanceId}: ${input.description}`);
118+
logger.verbose(`Message: ${input.message}`);
119+
}
120+
},
121+
122+
logReturns: (result, { logger }) => {
123+
logger.verbose(`Received ${result.response.length} characters in response`);
124+
},
125+
};

‎src/tools/interaction/subAgentStart.test.ts‎

Lines changed: 223 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -0,0 +1,223 @@
1+
import { describe, it, expect, vi, beforeEach, afterEach } from "vitest";
2+
import { subAgentStartTool } from "./subAgentStart.js";
3+
import { subAgentMessageTool } from "./subAgentMessage.js";
4+
import { Logger } from "../../utils/logger.js";
5+
import { subAgentStates } from "./subAgentState.js";
6+
7+
const logger = new Logger({ name: "subAgentStart", logLevel: "warn" });
8+
9+
// Mock Anthropic client response
10+
const mockResponse = {
11+
content: [
12+
{
13+
type: "text",
14+
text: "Initial response from sub-agent",
15+
},
16+
],
17+
usage: { input_tokens: 10, output_tokens: 10 },
18+
};
19+
20+
// Mock message response
21+
const mockMessageResponse = {
22+
content: [
23+
{
24+
type: "text",
25+
text: "Response to message",
26+
},
27+
],
28+
usage: { input_tokens: 5, output_tokens: 5 },
29+
};
30+
31+
// Mock message creation function
32+
const mockCreateMessage = vi.fn();
33+
34+
// Mock Anthropic SDK
35+
vi.mock("@anthropic-ai/sdk", () => ({
36+
default: class {
37+
messages = {
38+
create: mockCreateMessage,
39+
};
40+
},
41+
}));
42+
43+
describe("subAgentStart and subAgentMessage", () => {
44+
beforeEach(() => {
45+
process.env.ANTHROPIC_API_KEY = "test-key";
46+
subAgentStates.clear();
47+
vi.clearAllMocks();
48+
// Reset mock implementation for each test
49+
mockCreateMessage.mockReset();
50+
});
51+
52+
afterEach(() => {
53+
vi.clearAllMocks();
54+
subAgentStates.clear();
55+
});
56+
57+
describe("subAgentStart", () => {
58+
it("should start a sub-agent and return an instance ID", async () => {
59+
mockCreateMessage.mockResolvedValueOnce(mockResponse);
60+
const result = await subAgentStartTool.execute(
61+
{
62+
prompt: "Test sub-agent task",
63+
description: "A test agent for unit testing",
64+
},
65+
{ logger },
66+
);
67+
68+
expect(result.instanceId).toBeDefined();
69+
expect(result.response).toContain("Initial response from sub-agent");
70+
expect(subAgentStates.has(result.instanceId)).toBe(true);
71+
});
72+
73+
it("should handle multiple concurrent sub-agents", async () => {
74+
mockCreateMessage.mockResolvedValueOnce(mockResponse)
75+
.mockResolvedValueOnce(mockResponse);
76+
const result1 = await subAgentStartTool.execute(
77+
{
78+
prompt: "First agent task",
79+
description: "First test agent",
80+
},
81+
{ logger },
82+
);
83+
84+
const result2 = await subAgentStartTool.execute(
85+
{
86+
prompt: "Second agent task",
87+
description: "Second test agent",
88+
},
89+
{ logger },
90+
);
91+
92+
expect(result1.instanceId).not.toBe(result2.instanceId);
93+
expect(subAgentStates.has(result1.instanceId)).toBe(true);
94+
expect(subAgentStates.has(result2.instanceId)).toBe(true);
95+
});
96+
97+
it("should validate description length", async () => {
98+
const longDescription =
99+
"This is a very long description that exceeds the maximum allowed length of 80 characters and should fail";
100+
101+
await expect(
102+
subAgentStartTool.execute(
103+
{
104+
prompt: "Test task",
105+
description: longDescription,
106+
},
107+
{ logger },
108+
),
109+
).rejects.toThrow();
110+
});
111+
112+
it("should handle API errors gracefully", async () => {
113+
delete process.env.ANTHROPIC_API_KEY;
114+
115+
await expect(
116+
subAgentStartTool.execute(
117+
{
118+
prompt: "Test task",
119+
description: "Should fail",
120+
},
121+
{ logger },
122+
),
123+
).rejects.toThrow("ANTHROPIC_API_KEY environment variable is not set");
124+
});
125+
});
126+
127+
describe("subAgentMessage", () => {
128+
let instanceId: string;
129+
130+
beforeEach(async () => {
131+
mockCreateMessage.mockResolvedValueOnce(mockResponse);
132+
const result = await subAgentStartTool.execute(
133+
{
134+
prompt: "Test sub-agent task",
135+
description: "Test agent for messaging",
136+
},
137+
{ logger },
138+
);
139+
instanceId = result.instanceId;
140+
});
141+
142+
it("should send a message to an existing sub-agent", async () => {
143+
mockCreateMessage.mockResolvedValueOnce(mockMessageResponse);
144+
const result = await subAgentMessageTool.execute(
145+
{
146+
instanceId,
147+
message: "Test message",
148+
description: "Test message interaction",
149+
},
150+
{ logger },
151+
);
152+
153+
expect(result.response).toContain("Response to message");
154+
});
155+
156+
it("should handle invalid instance IDs", async () => {
157+
await expect(
158+
subAgentMessageTool.execute(
159+
{
160+
instanceId: "invalid-id",
161+
message: "Test message",
162+
description: "Should fail",
163+
},
164+
{ logger },
165+
),
166+
).rejects.toThrow("Sub-agent not found");
167+
});
168+
169+
it("should handle messaging aborted agents", async () => {
170+
// Abort the agent
171+
await subAgentMessageTool.execute(
172+
{
173+
instanceId,
174+
description: "Aborting agent",
175+
abort: true,
176+
},
177+
{ logger },
178+
);
179+
180+
// Try to send a message after abort
181+
await expect(
182+
subAgentMessageTool.execute(
183+
{
184+
instanceId,
185+
message: "Test message",
186+
description: "Should fail - agent aborted",
187+
},
188+
{ logger },
189+
),
190+
).rejects.toThrow("Sub-agent has been aborted");
191+
});
192+
193+
it("should abort a sub-agent", async () => {
194+
const result = await subAgentMessageTool.execute(
195+
{
196+
instanceId,
197+
description: "Aborting agent",
198+
abort: true,
199+
},
200+
{ logger },
201+
);
202+
203+
expect(result.response).toContain("Sub-agent aborted");
204+
expect(subAgentStates.get(instanceId)?.aborted).toBe(true);
205+
});
206+
207+
it("should validate message description length", async () => {
208+
const longDescription =
209+
"This is a very long description that exceeds the maximum allowed length of 80 characters and should fail";
210+
211+
await expect(
212+
subAgentMessageTool.execute(
213+
{
214+
instanceId,
215+
message: "Test message",
216+
description: longDescription,
217+
},
218+
{ logger },
219+
),
220+
).rejects.toThrow();
221+
});
222+
});
223+
});

0 commit comments

Comments
 (0)