From 1f9e5b489547d4ab1d7cdd729c47b8ca4c25b38b Mon Sep 17 00:00:00 2001 From: k2v7n24cbf-cyber Date: Mon, 20 Jul 2026 11:30:19 +0300 Subject: [PATCH] Stream OpenAI OAuth structured helpers --- packages/agent-core/src/subagent.ts | 97 +++++++++++++++++------ packages/agent-core/test/subagent.test.ts | 52 +++++++++++- 2 files changed, 123 insertions(+), 26 deletions(-) diff --git a/packages/agent-core/src/subagent.ts b/packages/agent-core/src/subagent.ts index f03a02d9..ea7a76db 100644 --- a/packages/agent-core/src/subagent.ts +++ b/packages/agent-core/src/subagent.ts @@ -13,6 +13,8 @@ import { generateObject, readUIMessageStream, stepCountIs, + streamObject, + type LanguageModelUsage, type StopCondition, type ToolSet, type UIMessage, @@ -231,19 +233,53 @@ async function runStructured( const startedAt = Date.now(); try { const subagentModel = resolveSubagentModel(args.parent.config.model); - const result = await generateObject({ - model: args.parent.resolveModel(subagentModel), - system: args.system, - schema: args.schema, - messages: [{ role: "user", content: args.user }], - maxOutputTokens: args.maxOutputTokens, - abortSignal: combined, - experimental_telemetry: { - isEnabled: true, - functionId: `${args.parent.config.id}:subagent.structured`, - }, - }); - if (args.usage && result.usage) { + let object: z.infer; + let usage: LanguageModelUsage | undefined; + + if (usesStreamingStructuredOutput(subagentModel)) { + const result = streamObject({ + model: args.parent.resolveModel(subagentModel), + system: args.system, + schema: args.schema, + messages: [{ role: "user", content: args.user }], + maxOutputTokens: args.maxOutputTokens, + abortSignal: combined, + experimental_telemetry: { + isEnabled: true, + functionId: `${args.parent.config.id}:subagent.structured`, + }, + }); + const streamFinished = drainStream(result.fullStream); + try { + object = (await result.object) as z.infer; + } finally { + await streamFinished.catch(() => { + // The object promise carries the meaningful failure for callers. + }); + } + try { + usage = await result.usage; + } catch { + // Structured result is authoritative; usage is best-effort metadata. + } + } else { + const result = await generateObject({ + model: args.parent.resolveModel(subagentModel), + system: args.system, + schema: args.schema, + messages: [{ role: "user", content: args.user }], + maxOutputTokens: args.maxOutputTokens, + abortSignal: combined, + experimental_telemetry: { + isEnabled: true, + functionId: `${args.parent.config.id}:subagent.structured`, + }, + }); + object = result.object as z.infer; + usage = result.usage; + } + + if (args.usage && usage) { args.parent.reportUsage({ agentId: args.parent.config.id, sessionId: args.usage.sessionId, @@ -251,12 +287,12 @@ async function runStructured( model: subagentModel, taskId: args.usage.taskId, tokens: { - inputTokens: result.usage.inputTokens, - outputTokens: result.usage.outputTokens, - totalTokens: result.usage.totalTokens, - cachedInputTokens: result.usage.inputTokenDetails?.cacheReadTokens, - cacheWriteTokens: result.usage.inputTokenDetails?.cacheWriteTokens, - reasoningTokens: result.usage.outputTokenDetails?.reasoningTokens, + inputTokens: usage.inputTokens, + outputTokens: usage.outputTokens, + totalTokens: usage.totalTokens, + cachedInputTokens: usage.inputTokenDetails?.cacheReadTokens, + cacheWriteTokens: usage.inputTokenDetails?.cacheWriteTokens, + reasoningTokens: usage.outputTokenDetails?.reasoningTokens, }, steps: 1, durationMs: Date.now() - startedAt, @@ -265,12 +301,12 @@ async function runStructured( return { mode: "structured", status: "completed", - object: result.object as z.infer, - usage: result.usage + object, + usage: usage ? { - inputTokens: result.usage.inputTokens, - outputTokens: result.usage.outputTokens, - totalTokens: result.usage.totalTokens, + inputTokens: usage.inputTokens, + outputTokens: usage.outputTokens, + totalTokens: usage.totalTokens, } : undefined, }; @@ -290,3 +326,16 @@ async function runStructured( }; } } + +function usesStreamingStructuredOutput(model: { + provider?: string; + auth?: string; +}): boolean { + return model.provider === "openai" && model.auth === "oauth"; +} + +async function drainStream(stream: AsyncIterable): Promise { + for await (const _ of stream) { + // Draining drives the SDK stream pipeline so final object/usage promises resolve. + } +} diff --git a/packages/agent-core/test/subagent.test.ts b/packages/agent-core/test/subagent.test.ts index 3dd962e3..55f7fffb 100644 --- a/packages/agent-core/test/subagent.test.ts +++ b/packages/agent-core/test/subagent.test.ts @@ -13,7 +13,7 @@ import { import { MemoryStore } from "@openacme/memory"; import { TaskStore } from "@openacme/tasks"; import type { ToolRegistry } from "@openacme/tools"; -import { MockLanguageModelV3 } from "ai/test"; +import { MockLanguageModelV3, simulateReadableStream } from "ai/test"; import { Agent } from "../src/agent.js"; import type { AgentConfig } from "../src/types.js"; import { runSubagent } from "../src/subagent.js"; @@ -41,7 +41,7 @@ function freshDb() { return db; } -function makeAgent(): Agent { +function makeAgent(model?: Partial): Agent { const tmpRoot = fs.mkdtempSync(path.join(os.tmpdir(), "openacme-subagent-")); const db = freshDb(); const sessionStore = createSessionStore(db); @@ -54,6 +54,7 @@ function makeAgent(): Agent { model: "test", apiKey: "x", auth: "api_key", + ...model, }, persona: "test", tools: [], @@ -307,6 +308,32 @@ function modelReturning(obj: unknown): MockLanguageModelV3 { }); } +function streamingModelReturning(obj: unknown): MockLanguageModelV3 { + return new MockLanguageModelV3({ + doGenerate: async () => { + throw new Error("doGenerate should not be used"); + }, + doStream: async () => { + const id = "txt-1"; + return { + stream: simulateReadableStream({ + chunks: [ + { type: "stream-start", warnings: [] }, + { type: "text-start", id }, + { type: "text-delta", id, delta: JSON.stringify(obj) }, + { type: "text-end", id }, + { + type: "finish", + finishReason: "stop", + usage: { inputTokens: 1, outputTokens: 1, totalTokens: 2 }, + }, + ], + }), + }; + }, + }); +} + describe("runSubagent (structured mode)", () => { beforeEach(() => { vi.restoreAllMocks(); @@ -330,6 +357,27 @@ describe("runSubagent (structured mode)", () => { } }); + it("streams structured output for OpenAI OAuth", async () => { + const agent = makeAgent({ auth: "oauth" }); + const model = streamingModelReturning({ selected: ["oauth"] }); + getModelMock.mockReturnValue(model); + + const out = await runSubagent({ + mode: "structured", + parent: agent, + system: "system", + user: "user", + schema: PickSchema, + }); + + expect(out.status).toBe("completed"); + if (out.mode === "structured") { + expect(out.object).toEqual({ selected: ["oauth"] }); + } + expect(model.doStreamCalls.length).toBe(1); + expect(model.doGenerateCalls.length).toBe(0); + }); + it("returns null + failed on schema mismatch", async () => { const agent = makeAgent(); getModelMock.mockReturnValue(