Skip to content

Commit 864c3d8

Browse files
charlesliencharleslien
committed
add STEP_TEXT to handleSteps
1 parent a05dcbd commit 864c3d8

10 files changed

Lines changed: 79 additions & 14 deletions

File tree

‎.agents/types/agent-definition.ts‎

Lines changed: 3 additions & 1 deletion
Original file line numberDiff line numberDiff line change
@@ -197,7 +197,7 @@ export interface AgentDefinition {
197197
* }
198198
*/
199199
handleSteps?: (context: AgentStepContext) => Generator<
200-
ToolCall | 'STEP' | 'STEP_ALL',
200+
ToolCall | 'STEP' | 'STEP_ALL' | StepText,
201201
void,
202202
{
203203
agentState: AgentState
@@ -233,6 +233,8 @@ export interface AgentStepContext {
233233
logger: Logger
234234
}
235235

236+
export type StepText = { type: 'STEP_TEXT'; text: string }
237+
236238
/**
237239
* Tool call object for handleSteps generator
238240
*/

‎common/src/templates/initial-agents-dir/types/agent-definition.ts‎

Lines changed: 3 additions & 1 deletion
Original file line numberDiff line numberDiff line change
@@ -197,7 +197,7 @@ export interface AgentDefinition {
197197
* }
198198
*/
199199
handleSteps?: (context: AgentStepContext) => Generator<
200-
ToolCall | 'STEP' | 'STEP_ALL',
200+
ToolCall | 'STEP' | 'STEP_ALL' | StepText,
201201
void,
202202
{
203203
agentState: AgentState
@@ -233,6 +233,8 @@ export interface AgentStepContext {
233233
logger: Logger
234234
}
235235

236+
export type StepText = { type: 'STEP_TEXT'; text: string }
237+
236238
/**
237239
* Tool call object for handleSteps generator
238240
*/

‎common/src/types/agent-template.ts‎

Lines changed: 3 additions & 1 deletion
Original file line numberDiff line numberDiff line change
@@ -82,8 +82,10 @@ export type AgentTemplate<
8282
handleSteps?: StepHandler<P, T> | string // Function or string of the generator code for running in a sandbox
8383
}
8484

85+
export type StepText = { type: 'STEP_TEXT'; text: string }
86+
8587
export type StepGenerator = Generator<
86-
Omit<ToolCall, 'toolCallId'> | 'STEP' | 'STEP_ALL', // Generic tool call type
88+
Omit<ToolCall, 'toolCallId'> | 'STEP' | 'STEP_ALL' | StepText, // Generic tool call type
8789
void,
8890
{
8991
agentState: PublicAgentState

‎evals/scaffolding.ts‎

Lines changed: 1 addition & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -205,6 +205,7 @@ export async function runAgentStepScaffolding(
205205
const result = await runAgentStep({
206206
...EVALS_AGENT_RUNTIME_IMPL,
207207
...agentRuntimeScopedImpl,
208+
textOverride: null,
208209
runId: 'test-run-id',
209210
userId: TEST_USER_ID,
210211
userInputId: generateCompactId(),

‎packages/agent-runtime/src/__tests__/read-docs-tool.test.ts‎

Lines changed: 6 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -151,6 +151,7 @@ describe('read_docs tool with researcher agent', () => {
151151

152152
const { agentState: newAgentState } = await runAgentStep({
153153
...agentRuntimeImpl,
154+
textOverride: null,
154155
runId: 'test-run-id',
155156
repoId: undefined,
156157
repoUrl: undefined,
@@ -227,6 +228,7 @@ describe('read_docs tool with researcher agent', () => {
227228

228229
await runAgentStep({
229230
...agentRuntimeImpl,
231+
textOverride: null,
230232
runId: 'test-run-id',
231233
repoId: undefined,
232234
repoUrl: undefined,
@@ -280,6 +282,7 @@ describe('read_docs tool with researcher agent', () => {
280282

281283
const { agentState: newAgentState } = await runAgentStep({
282284
...agentRuntimeImpl,
285+
textOverride: null,
283286
runId: 'test-run-id',
284287
repoId: undefined,
285288
repoUrl: undefined,
@@ -349,6 +352,7 @@ describe('read_docs tool with researcher agent', () => {
349352

350353
const { agentState: newAgentState } = await runAgentStep({
351354
...agentRuntimeImpl,
355+
textOverride: null,
352356
runId: 'test-run-id',
353357
repoId: undefined,
354358
repoUrl: undefined,
@@ -417,6 +421,7 @@ describe('read_docs tool with researcher agent', () => {
417421

418422
const { agentState: newAgentState } = await runAgentStep({
419423
...agentRuntimeImpl,
424+
textOverride: null,
420425
runId: 'test-run-id',
421426
repoId: undefined,
422427
repoUrl: undefined,
@@ -486,6 +491,7 @@ describe('read_docs tool with researcher agent', () => {
486491

487492
const { agentState: newAgentState } = await runAgentStep({
488493
...agentRuntimeImpl,
494+
textOverride: null,
489495
runId: 'test-run-id',
490496
repoId: undefined,
491497
repoUrl: undefined,

‎packages/agent-runtime/src/__tests__/run-agent-step-tools.test.ts‎

Lines changed: 6 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -163,6 +163,7 @@ describe('runAgentStep - set_output tool', () => {
163163

164164
const result = await runAgentStep({
165165
...agentRuntimeImpl,
166+
textOverride: null,
166167
runId: 'test-run-id',
167168
repoId: undefined,
168169
repoUrl: undefined,
@@ -207,6 +208,7 @@ describe('runAgentStep - set_output tool', () => {
207208

208209
const result = await runAgentStep({
209210
...agentRuntimeImpl,
211+
textOverride: null,
210212
runId: 'test-run-id',
211213
repoId: undefined,
212214
repoUrl: undefined,
@@ -257,6 +259,7 @@ describe('runAgentStep - set_output tool', () => {
257259

258260
const result = await runAgentStep({
259261
...agentRuntimeImpl,
262+
textOverride: null,
260263
runId: 'test-run-id',
261264
repoId: undefined,
262265
repoUrl: undefined,
@@ -298,6 +301,7 @@ describe('runAgentStep - set_output tool', () => {
298301

299302
const result = await runAgentStep({
300303
...agentRuntimeImpl,
304+
textOverride: null,
301305
runId: 'test-run-id',
302306
repoId: undefined,
303307
repoUrl: undefined,
@@ -394,6 +398,7 @@ describe('runAgentStep - set_output tool', () => {
394398

395399
const result = await runAgentStep({
396400
...agentRuntimeImpl,
401+
textOverride: null,
397402
runId: 'test-run-id',
398403
repoId: undefined,
399404
repoUrl: undefined,
@@ -556,6 +561,7 @@ describe('runAgentStep - set_output tool', () => {
556561

557562
const result = await runAgentStep({
558563
...agentRuntimeImpl,
564+
textOverride: null,
559565
runId: 'test-run-id',
560566
repoId: undefined,
561567
repoUrl: undefined,

‎packages/agent-runtime/src/__tests__/web-search-tool.test.ts‎

Lines changed: 8 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -124,6 +124,7 @@ describe('web_search tool with researcher agent', () => {
124124

125125
await runAgentStep({
126126
...agentRuntimeImpl,
127+
textOverride: null,
127128
system: 'Test system prompt',
128129
userId: TEST_USER_ID,
129130
userInputId: 'test-input',
@@ -177,6 +178,7 @@ describe('web_search tool with researcher agent', () => {
177178

178179
const { agentState: newAgentState } = await runAgentStep({
179180
...agentRuntimeImpl,
181+
textOverride: null,
180182
system: 'Test system prompt',
181183
userId: TEST_USER_ID,
182184
userInputId: 'test-input',
@@ -239,6 +241,7 @@ describe('web_search tool with researcher agent', () => {
239241

240242
await runAgentStep({
241243
...agentRuntimeImpl,
244+
textOverride: null,
242245
system: 'Test system prompt',
243246
userId: TEST_USER_ID,
244247
userInputId: 'test-input',
@@ -286,6 +289,7 @@ describe('web_search tool with researcher agent', () => {
286289

287290
const { agentState: newAgentState } = await runAgentStep({
288291
...agentRuntimeImpl,
292+
textOverride: null,
289293
system: 'Test system prompt',
290294
userId: TEST_USER_ID,
291295
userInputId: 'test-input',
@@ -347,6 +351,7 @@ describe('web_search tool with researcher agent', () => {
347351

348352
const { agentState: newAgentState } = await runAgentStep({
349353
...agentRuntimeImpl,
354+
textOverride: null,
350355
system: 'Test system prompt',
351356
userId: TEST_USER_ID,
352357
userInputId: 'test-input',
@@ -407,6 +412,7 @@ describe('web_search tool with researcher agent', () => {
407412

408413
const { agentState: newAgentState } = await runAgentStep({
409414
...agentRuntimeImpl,
415+
textOverride: null,
410416
system: 'Test system prompt',
411417
userId: TEST_USER_ID,
412418
userInputId: 'test-input',
@@ -457,6 +463,7 @@ describe('web_search tool with researcher agent', () => {
457463

458464
const { agentState: newAgentState } = await runAgentStep({
459465
...agentRuntimeImpl,
466+
textOverride: null,
460467
system: 'Test system prompt',
461468
userId: TEST_USER_ID,
462469
userInputId: 'test-input',
@@ -519,6 +526,7 @@ describe('web_search tool with researcher agent', () => {
519526

520527
const { agentState: newAgentState } = await runAgentStep({
521528
...agentRuntimeImpl,
529+
textOverride: null,
522530
system: 'Test system prompt',
523531
userId: TEST_USER_ID,
524532
userInputId: 'test-input',

‎packages/agent-runtime/src/prompt-agent-stream.ts‎

Lines changed: 13 additions & 2 deletions
Original file line numberDiff line numberDiff line change
@@ -23,6 +23,7 @@ export const getAgentStreamFromTemplate = (params: {
2323
onCostCalculated?: (credits: number) => Promise<void>
2424
agentId?: string
2525
includeCacheControl?: boolean
26+
textOverride: string | null
2627

2728
template: AgentTemplate
2829
logger: Logger
@@ -31,7 +32,7 @@ export const getAgentStreamFromTemplate = (params: {
3132
liveUserInputRecord: UserInputRecord
3233
sessionConnections: SessionRecord
3334
trackEvent: TrackEventFn
34-
}) => {
35+
}): { getStream: (messages: Message[]) => ReturnType<PromptAiSdkStreamFn> } => {
3536
const {
3637
apiKey,
3738
runId,
@@ -42,6 +43,7 @@ export const getAgentStreamFromTemplate = (params: {
4243
onCostCalculated,
4344
agentId,
4445
includeCacheControl,
46+
textOverride,
4547
template,
4648
logger,
4749
sendAction,
@@ -51,13 +53,22 @@ export const getAgentStreamFromTemplate = (params: {
5153
trackEvent,
5254
} = params
5355

56+
if (textOverride !== null) {
57+
return {
58+
getStream: async function* stream(): ReturnType<PromptAiSdkStreamFn> {
59+
yield { type: 'text', text: textOverride!, agentId }
60+
return null
61+
},
62+
}
63+
}
64+
5465
if (!template) {
5566
throw new Error('Agent template is null/undefined')
5667
}
5768

5869
const { model } = template
5970

60-
const getStream = (messages: Message[]) => {
71+
const getStream = (messages: Message[]): ReturnType<PromptAiSdkStreamFn> => {
6172
const aiSdkStreamParams: ParamsOf<PromptAiSdkStreamFn> = {
6273
apiKey,
6374
runId,

‎packages/agent-runtime/src/run-agent-step.ts‎

Lines changed: 16 additions & 6 deletions
Original file line numberDiff line numberDiff line change
@@ -472,7 +472,12 @@ export async function loopAgentSteps(
472472
> &
473473
ParamsExcluding<
474474
typeof runAgentStep,
475-
'agentState' | 'prompt' | 'spawnParams' | 'system' | 'runId'
475+
| 'agentState'
476+
| 'prompt'
477+
| 'spawnParams'
478+
| 'system'
479+
| 'runId'
480+
| 'textOverride'
476481
> &
477482
ParamsExcluding<
478483
AddAgentStepFn,
@@ -644,12 +649,9 @@ export async function loopAgentSteps(
644649
const startTime = new Date()
645650

646651
// 1. Run programmatic step first if it exists
652+
let textOverride = null
647653
if (agentTemplate.handleSteps) {
648-
const {
649-
agentState: programmaticAgentState,
650-
endTurn,
651-
stepNumber,
652-
} = await runProgrammaticStep({
654+
const programmaticResult = await runProgrammaticStep({
653655
...params,
654656
runId,
655657
agentState: currentAgentState,
@@ -661,6 +663,13 @@ export async function loopAgentSteps(
661663
stepsComplete: shouldEndTurn,
662664
stepNumber: totalSteps,
663665
})
666+
const {
667+
agentState: programmaticAgentState,
668+
endTurn,
669+
stepNumber,
670+
} = programmaticResult
671+
textOverride = programmaticResult.textOverride
672+
664673
currentAgentState = programmaticAgentState
665674
totalSteps = stepNumber
666675

@@ -715,6 +724,7 @@ export async function loopAgentSteps(
715724
messageId,
716725
} = await runAgentStep({
717726
...params,
727+
textOverride: textOverride,
718728
runId,
719729
agentState: currentAgentState,
720730
prompt: currentPrompt,

‎packages/agent-runtime/src/run-programmatic-step.ts‎

Lines changed: 20 additions & 3 deletions
Original file line numberDiff line numberDiff line change
@@ -93,7 +93,12 @@ export async function runProgrammaticStep(
9393
| 'startTime'
9494
| 'messageId'
9595
>,
96-
): Promise<{ agentState: AgentState; endTurn: boolean; stepNumber: number }> {
96+
): Promise<{
97+
agentState: AgentState
98+
textOverride: string | null
99+
endTurn: boolean
100+
stepNumber: number
101+
}> {
97102
const {
98103
agentState,
99104
template,
@@ -181,7 +186,7 @@ export async function runProgrammaticStep(
181186
// Clear the STEP_ALL mode. Stepping can continue if handleSteps doesn't return.
182187
runIdToStepAll.delete(agentState.runId)
183188
} else {
184-
return { agentState, endTurn: false, stepNumber }
189+
return { agentState, textOverride: null, endTurn: false, stepNumber }
185190
}
186191
}
187192

@@ -221,6 +226,7 @@ export async function runProgrammaticStep(
221226

222227
let toolResult: ToolResultOutput[] = []
223228
let endTurn = false
229+
let textOverride: string | null = null
224230

225231
let startTime = new Date()
226232
let creditsBefore = agentState.directCreditsUsed
@@ -257,6 +263,11 @@ export async function runProgrammaticStep(
257263
break
258264
}
259265

266+
if (result.value.type === 'STEP_TEXT') {
267+
textOverride = result.value.text
268+
break
269+
}
270+
260271
// Process tool calls yielded by the generator
261272
const toolCallWithoutId = result.value
262273
const toolCall = {
@@ -396,7 +407,12 @@ export async function runProgrammaticStep(
396407
}
397408
} while (true)
398409

399-
return { agentState: state.agentState, endTurn, stepNumber }
410+
return {
411+
agentState: state.agentState,
412+
textOverride: textOverride,
413+
endTurn,
414+
stepNumber,
415+
}
400416
} catch (error) {
401417
endTurn = true
402418

@@ -442,6 +458,7 @@ export async function runProgrammaticStep(
442458

443459
return {
444460
agentState: state.agentState,
461+
textOverride: null,
445462
endTurn,
446463
stepNumber,
447464
}

0 commit comments

Comments
 (0)