refactor: unify structured generation pipeline
This commit is contained in:
Mohamed Boudra
@@ -1,5 +1,9 @@
|
||||
import type { Command } from 'commander'
|
||||
import type { AgentSnapshotPayload } from '@getpaseo/server'
|
||||
import {
|
||||
getStructuredAgentResponse,
|
||||
StructuredAgentResponseError,
|
||||
type AgentSnapshotPayload,
|
||||
} from '@getpaseo/server'
|
||||
import { connectToDaemon, getDaemonHost } from '../../utils/client.js'
|
||||
import type { CommandOptions, SingleResult, OutputSchema, CommandError } from '../../output/index.js'
|
||||
import { readFileSync } from 'node:fs'
|
||||
@@ -105,98 +109,55 @@ function loadOutputSchema(value: string): Record<string, unknown> {
|
||||
return parsed as Record<string, unknown>
|
||||
}
|
||||
|
||||
function extractFirstJsonObject(text: string): string | null {
|
||||
const source = text.trim()
|
||||
if (!source) {
|
||||
return null
|
||||
class StructuredRunStatusError extends Error {
|
||||
readonly kind: 'timeout' | 'permission' | 'error' | 'empty'
|
||||
|
||||
constructor(kind: 'timeout' | 'permission' | 'error' | 'empty', message: string) {
|
||||
super(message)
|
||||
this.name = 'StructuredRunStatusError'
|
||||
this.kind = kind
|
||||
}
|
||||
}
|
||||
|
||||
type ConnectedDaemonClient = Awaited<ReturnType<typeof connectToDaemon>>
|
||||
|
||||
export interface StructuredResponseTimelineClient {
|
||||
fetchAgentTimeline: ConnectedDaemonClient['fetchAgentTimeline']
|
||||
}
|
||||
|
||||
export async function resolveStructuredResponseMessage(options: {
|
||||
client: StructuredResponseTimelineClient
|
||||
agentId: string
|
||||
lastMessage: string | null
|
||||
}): Promise<string | null> {
|
||||
const direct = options.lastMessage?.trim()
|
||||
if (direct) {
|
||||
return direct
|
||||
}
|
||||
|
||||
const startIndexes: number[] = []
|
||||
for (let i = 0; i < source.length; i += 1) {
|
||||
if (source[i] === '{') {
|
||||
startIndexes.push(i)
|
||||
}
|
||||
}
|
||||
|
||||
for (const start of startIndexes) {
|
||||
let depth = 0
|
||||
let inString = false
|
||||
let escaped = false
|
||||
|
||||
for (let i = start; i < source.length; i += 1) {
|
||||
const ch = source[i]!
|
||||
|
||||
if (inString) {
|
||||
if (escaped) {
|
||||
escaped = false
|
||||
continue
|
||||
}
|
||||
if (ch === '\\') {
|
||||
escaped = true
|
||||
continue
|
||||
}
|
||||
if (ch === '"') {
|
||||
inString = false
|
||||
}
|
||||
try {
|
||||
const timeline = await options.client.fetchAgentTimeline(options.agentId, {
|
||||
direction: 'tail',
|
||||
projection: 'projected',
|
||||
limit: 200,
|
||||
})
|
||||
for (let index = timeline.entries.length - 1; index >= 0; index -= 1) {
|
||||
const entry = timeline.entries[index]
|
||||
if (!entry || entry.item.type !== 'assistant_message') {
|
||||
continue
|
||||
}
|
||||
|
||||
if (ch === '"') {
|
||||
inString = true
|
||||
continue
|
||||
}
|
||||
|
||||
if (ch === '{') {
|
||||
depth += 1
|
||||
continue
|
||||
}
|
||||
if (ch === '}') {
|
||||
depth -= 1
|
||||
if (depth === 0) {
|
||||
const candidate = source.slice(start, i + 1).trim()
|
||||
try {
|
||||
JSON.parse(candidate)
|
||||
return candidate
|
||||
} catch {
|
||||
// Keep scanning.
|
||||
}
|
||||
}
|
||||
const text = entry.item.text.trim()
|
||||
if (text.length > 0) {
|
||||
return text
|
||||
}
|
||||
}
|
||||
} catch {
|
||||
// Leave empty; caller will surface a consistent structured-output failure message.
|
||||
}
|
||||
|
||||
return null
|
||||
}
|
||||
|
||||
function parseStructuredOutput(lastMessage: string): Record<string, unknown> {
|
||||
const trimmed = lastMessage.trim()
|
||||
const fenced = trimmed.match(/```(?:json)?\s*\n([\s\S]*?)\n```/)
|
||||
const jsonText = fenced?.[1]?.trim() ?? extractFirstJsonObject(trimmed) ?? trimmed
|
||||
|
||||
let parsed: unknown
|
||||
try {
|
||||
parsed = JSON.parse(jsonText)
|
||||
} catch (err) {
|
||||
const message = err instanceof Error ? err.message : String(err)
|
||||
const error: CommandError = {
|
||||
code: 'OUTPUT_SCHEMA_FAILED',
|
||||
message: 'Agent response is not valid JSON',
|
||||
details: message,
|
||||
}
|
||||
throw error
|
||||
}
|
||||
|
||||
if (!parsed || typeof parsed !== 'object' || Array.isArray(parsed)) {
|
||||
const error: CommandError = {
|
||||
code: 'OUTPUT_SCHEMA_FAILED',
|
||||
message: 'Agent response JSON must be an object',
|
||||
}
|
||||
throw error
|
||||
}
|
||||
|
||||
return parsed as Record<string, unknown>
|
||||
}
|
||||
|
||||
function structuredRunSchema(output: Record<string, unknown>): OutputSchema<AgentRunResult> {
|
||||
return {
|
||||
...agentRunSchema,
|
||||
@@ -318,6 +279,107 @@ export async function runRunCommand(
|
||||
labels['ui'] = 'true'
|
||||
}
|
||||
|
||||
if (outputSchema) {
|
||||
let structuredAgent: AgentSnapshotPayload | null = null
|
||||
|
||||
const callStructuredTurn = async (structuredPrompt: string): Promise<string> => {
|
||||
if (!structuredAgent) {
|
||||
structuredAgent = await client.createAgent({
|
||||
provider: (options.provider as 'claude' | 'codex' | 'opencode') ?? 'claude',
|
||||
cwd,
|
||||
title: options.name,
|
||||
modeId: options.mode,
|
||||
model: options.model,
|
||||
initialPrompt: structuredPrompt,
|
||||
images,
|
||||
git,
|
||||
worktreeName: options.worktree,
|
||||
labels: Object.keys(labels).length > 0 ? labels : undefined,
|
||||
})
|
||||
} else {
|
||||
await client.sendMessage(structuredAgent.id, structuredPrompt)
|
||||
}
|
||||
|
||||
const state = await client.waitForFinish(structuredAgent.id, 10 * 60 * 1000)
|
||||
if (state.status === 'timeout') {
|
||||
throw new StructuredRunStatusError('timeout', 'Timed out waiting for structured output')
|
||||
}
|
||||
if (state.status === 'permission') {
|
||||
throw new StructuredRunStatusError(
|
||||
'permission',
|
||||
'Agent is waiting for permission before producing structured output'
|
||||
)
|
||||
}
|
||||
if (state.status === 'error') {
|
||||
throw new StructuredRunStatusError(
|
||||
'error',
|
||||
state.error ?? 'Agent failed before producing structured output'
|
||||
)
|
||||
}
|
||||
|
||||
const lastMessage = await resolveStructuredResponseMessage({
|
||||
client,
|
||||
agentId: structuredAgent.id,
|
||||
lastMessage: state.lastMessage,
|
||||
})
|
||||
if (!lastMessage) {
|
||||
throw new StructuredRunStatusError(
|
||||
'empty',
|
||||
'Agent finished without a structured output message'
|
||||
)
|
||||
}
|
||||
|
||||
return lastMessage
|
||||
}
|
||||
|
||||
let output: Record<string, unknown>
|
||||
try {
|
||||
output = await getStructuredAgentResponse<Record<string, unknown>>({
|
||||
caller: callStructuredTurn,
|
||||
prompt,
|
||||
schema: outputSchema,
|
||||
schemaName: 'RunOutput',
|
||||
maxRetries: 2,
|
||||
})
|
||||
} catch (err) {
|
||||
if (err instanceof StructuredRunStatusError) {
|
||||
const error: CommandError = {
|
||||
code: 'OUTPUT_SCHEMA_FAILED',
|
||||
message: err.message,
|
||||
}
|
||||
throw error
|
||||
}
|
||||
if (err instanceof StructuredAgentResponseError) {
|
||||
const error: CommandError = {
|
||||
code: 'OUTPUT_SCHEMA_FAILED',
|
||||
message: 'Agent response did not match the required output schema',
|
||||
details:
|
||||
err.validationErrors.length > 0
|
||||
? err.validationErrors.join('\n')
|
||||
: err.lastResponse || 'No response',
|
||||
}
|
||||
throw error
|
||||
}
|
||||
throw err
|
||||
}
|
||||
|
||||
if (!structuredAgent) {
|
||||
const error: CommandError = {
|
||||
code: 'OUTPUT_SCHEMA_FAILED',
|
||||
message: 'Agent finished without a structured output message',
|
||||
}
|
||||
throw error
|
||||
}
|
||||
|
||||
await client.close()
|
||||
|
||||
return {
|
||||
type: 'single',
|
||||
data: toRunResult(structuredAgent, 'completed'),
|
||||
schema: structuredRunSchema(output),
|
||||
}
|
||||
}
|
||||
|
||||
// Create the agent
|
||||
const agent = await client.createAgent({
|
||||
provider: (options.provider as 'claude' | 'codex' | 'opencode') ?? 'claude',
|
||||
@@ -326,59 +388,12 @@ export async function runRunCommand(
|
||||
modeId: options.mode,
|
||||
model: options.model,
|
||||
initialPrompt: prompt,
|
||||
outputSchema,
|
||||
images,
|
||||
git,
|
||||
worktreeName: options.worktree,
|
||||
labels: Object.keys(labels).length > 0 ? labels : undefined,
|
||||
})
|
||||
|
||||
if (outputSchema) {
|
||||
const state = await client.waitForFinish(agent.id, 10 * 60 * 1000)
|
||||
|
||||
if (state.status === 'timeout') {
|
||||
const error: CommandError = {
|
||||
code: 'OUTPUT_SCHEMA_FAILED',
|
||||
message: 'Timed out waiting for structured output',
|
||||
}
|
||||
throw error
|
||||
}
|
||||
|
||||
if (state.status === 'permission') {
|
||||
const error: CommandError = {
|
||||
code: 'OUTPUT_SCHEMA_FAILED',
|
||||
message: 'Agent is waiting for permission before producing structured output',
|
||||
}
|
||||
throw error
|
||||
}
|
||||
|
||||
if (state.status === 'error') {
|
||||
const error: CommandError = {
|
||||
code: 'OUTPUT_SCHEMA_FAILED',
|
||||
message: state.error ?? 'Agent failed before producing structured output',
|
||||
}
|
||||
throw error
|
||||
}
|
||||
|
||||
const lastMessage = state.lastMessage?.trim()
|
||||
if (!lastMessage) {
|
||||
const error: CommandError = {
|
||||
code: 'OUTPUT_SCHEMA_FAILED',
|
||||
message: 'Agent finished without a structured output message',
|
||||
}
|
||||
throw error
|
||||
}
|
||||
|
||||
const output = parseStructuredOutput(lastMessage)
|
||||
await client.close()
|
||||
|
||||
return {
|
||||
type: 'single',
|
||||
data: toRunResult(agent, 'completed'),
|
||||
schema: structuredRunSchema(output),
|
||||
}
|
||||
}
|
||||
|
||||
// Default run behavior is foreground: wait for completion unless --detach is set.
|
||||
if (!options.detach) {
|
||||
const state = await client.waitForFinish(agent.id, 10 * 60 * 1000)
|
||||
|
||||
@@ -0,0 +1,97 @@
|
||||
#!/usr/bin/env npx tsx
|
||||
|
||||
import assert from 'node:assert'
|
||||
import {
|
||||
resolveStructuredResponseMessage,
|
||||
type StructuredResponseTimelineClient,
|
||||
} from '../src/commands/agent/run.ts'
|
||||
|
||||
type TimelineEntry = {
|
||||
item: {
|
||||
type: string
|
||||
text?: string
|
||||
}
|
||||
}
|
||||
|
||||
function createClient(options: {
|
||||
entries?: TimelineEntry[]
|
||||
throwOnFetch?: boolean
|
||||
onFetch?: () => void
|
||||
}): StructuredResponseTimelineClient {
|
||||
return {
|
||||
fetchAgentTimeline: async () => {
|
||||
options.onFetch?.()
|
||||
if (options.throwOnFetch) {
|
||||
throw new Error('timeline unavailable')
|
||||
}
|
||||
return {
|
||||
entries: options.entries ?? [],
|
||||
} as Awaited<ReturnType<StructuredResponseTimelineClient['fetchAgentTimeline']>>
|
||||
},
|
||||
}
|
||||
}
|
||||
|
||||
console.log('=== Run Output Schema Helper Tests ===\n')
|
||||
|
||||
// Test 1: Direct lastMessage is returned without timeline fetch.
|
||||
{
|
||||
let fetchCount = 0
|
||||
const client = createClient({ onFetch: () => fetchCount += 1 })
|
||||
const result = await resolveStructuredResponseMessage({
|
||||
client,
|
||||
agentId: 'agent-1',
|
||||
lastMessage: ' {"summary":"ok"} ',
|
||||
})
|
||||
assert.strictEqual(result, '{"summary":"ok"}')
|
||||
assert.strictEqual(fetchCount, 0, 'should not fetch timeline when lastMessage exists')
|
||||
console.log('✓ returns direct lastMessage when present')
|
||||
}
|
||||
|
||||
// Test 2: Falls back to latest assistant message from timeline.
|
||||
{
|
||||
const client = createClient({
|
||||
entries: [
|
||||
{ item: { type: 'user_message', text: 'prompt' } },
|
||||
{ item: { type: 'assistant_message', text: '' } },
|
||||
{ item: { type: 'assistant_message', text: ' {"summary":"from timeline"} ' } },
|
||||
],
|
||||
})
|
||||
const result = await resolveStructuredResponseMessage({
|
||||
client,
|
||||
agentId: 'agent-2',
|
||||
lastMessage: null,
|
||||
})
|
||||
assert.strictEqual(result, '{"summary":"from timeline"}')
|
||||
console.log('✓ falls back to latest assistant timeline entry')
|
||||
}
|
||||
|
||||
// Test 3: Returns null when timeline has no assistant messages.
|
||||
{
|
||||
const client = createClient({
|
||||
entries: [
|
||||
{ item: { type: 'user_message', text: 'prompt' } },
|
||||
{ item: { type: 'reasoning', text: 'thinking' } },
|
||||
],
|
||||
})
|
||||
const result = await resolveStructuredResponseMessage({
|
||||
client,
|
||||
agentId: 'agent-3',
|
||||
lastMessage: null,
|
||||
})
|
||||
assert.strictEqual(result, null)
|
||||
console.log('✓ returns null when no assistant messages exist')
|
||||
}
|
||||
|
||||
// Test 4: Returns null if timeline fetch fails.
|
||||
{
|
||||
const client = createClient({ throwOnFetch: true })
|
||||
const result = await resolveStructuredResponseMessage({
|
||||
client,
|
||||
agentId: 'agent-4',
|
||||
lastMessage: null,
|
||||
})
|
||||
assert.strictEqual(result, null)
|
||||
console.log('✓ returns null when timeline fetch throws')
|
||||
}
|
||||
|
||||
console.log('\n=== All helper tests passed ===')
|
||||
@@ -0,0 +1,139 @@
|
||||
#!/usr/bin/env npx tsx
|
||||
|
||||
import assert from 'node:assert'
|
||||
import { createE2ETestContext, type TestDaemonContext } from '../helpers/test-daemon.ts'
|
||||
|
||||
interface E2EContext extends TestDaemonContext {
|
||||
paseo: (args: string[], opts?: { timeout?: number; cwd?: string }) => Promise<{
|
||||
exitCode: number
|
||||
stdout: string
|
||||
stderr: string
|
||||
}>
|
||||
}
|
||||
|
||||
const schema = JSON.stringify({
|
||||
type: 'object',
|
||||
properties: {
|
||||
summary: { type: 'string' },
|
||||
},
|
||||
required: ['summary'],
|
||||
additionalProperties: false,
|
||||
})
|
||||
|
||||
const impossibleSchema = JSON.stringify({
|
||||
type: 'object',
|
||||
properties: {
|
||||
summary: { type: 'string' },
|
||||
},
|
||||
required: ['summary'],
|
||||
allOf: [
|
||||
{ properties: { summary: { type: 'string' } } },
|
||||
{ properties: { summary: { type: 'number' } } },
|
||||
],
|
||||
})
|
||||
|
||||
let ctx: E2EContext
|
||||
|
||||
async function setup(): Promise<void> {
|
||||
ctx = await createE2ETestContext({ timeout: 180000 })
|
||||
}
|
||||
|
||||
async function cleanup(): Promise<void> {
|
||||
if (ctx) {
|
||||
await ctx.stop()
|
||||
}
|
||||
}
|
||||
|
||||
async function runProviderCase(input: {
|
||||
provider: 'claude' | 'codex' | 'opencode'
|
||||
mode: string
|
||||
model: string
|
||||
}): Promise<void> {
|
||||
const result = await ctx.paseo(
|
||||
[
|
||||
'run',
|
||||
'--provider',
|
||||
input.provider,
|
||||
'--mode',
|
||||
input.mode,
|
||||
'--model',
|
||||
input.model,
|
||||
'--output-schema',
|
||||
schema,
|
||||
`Return valid JSON with a short summary for provider ${input.provider}.`,
|
||||
],
|
||||
{ timeout: 180000 }
|
||||
)
|
||||
|
||||
assert.strictEqual(
|
||||
result.exitCode,
|
||||
0,
|
||||
`expected ${input.provider} run to succeed, got ${result.exitCode}\nstdout:\n${result.stdout}\nstderr:\n${result.stderr}`
|
||||
)
|
||||
|
||||
const parsed = JSON.parse(result.stdout.trim()) as { summary?: unknown }
|
||||
assert.strictEqual(
|
||||
typeof parsed.summary,
|
||||
'string',
|
||||
`expected ${input.provider} output to contain string summary, got: ${result.stdout}`
|
||||
)
|
||||
assert(parsed.summary && parsed.summary.length > 0, 'summary must not be empty')
|
||||
}
|
||||
|
||||
async function test_all_providers_return_structured_output(): Promise<void> {
|
||||
await runProviderCase({
|
||||
provider: 'claude',
|
||||
mode: 'bypassPermissions',
|
||||
model: 'haiku',
|
||||
})
|
||||
await runProviderCase({
|
||||
provider: 'codex',
|
||||
mode: 'full-access',
|
||||
model: 'gpt-5.3-codex',
|
||||
})
|
||||
await runProviderCase({
|
||||
provider: 'opencode',
|
||||
mode: 'default',
|
||||
model: 'opencode/kimi-k2.5-free',
|
||||
})
|
||||
}
|
||||
|
||||
async function test_schema_validation_is_enforced(): Promise<void> {
|
||||
const result = await ctx.paseo(
|
||||
[
|
||||
'run',
|
||||
'--provider',
|
||||
'claude',
|
||||
'--mode',
|
||||
'bypassPermissions',
|
||||
'--model',
|
||||
'haiku',
|
||||
'--output-schema',
|
||||
impossibleSchema,
|
||||
'Return exactly {"summary":"ok"} and nothing else.',
|
||||
],
|
||||
{ timeout: 180000 }
|
||||
)
|
||||
|
||||
assert.notStrictEqual(result.exitCode, 0, 'expected impossible schema to fail')
|
||||
const output = `${result.stdout}\n${result.stderr}`
|
||||
assert(
|
||||
output.includes('OUTPUT_SCHEMA_FAILED'),
|
||||
`expected OUTPUT_SCHEMA_FAILED for impossible schema\nstdout:\n${result.stdout}\nstderr:\n${result.stderr}`
|
||||
)
|
||||
}
|
||||
|
||||
async function main(): Promise<void> {
|
||||
try {
|
||||
await setup()
|
||||
await test_all_providers_return_structured_output()
|
||||
await test_schema_validation_is_enforced()
|
||||
} catch (error) {
|
||||
console.error(error)
|
||||
process.exitCode = 1
|
||||
} finally {
|
||||
await cleanup()
|
||||
}
|
||||
}
|
||||
|
||||
main()
|
||||
Reference in New Issue
Block a user