fix(sdk): add Query.continue() to resume after a failed model turn - #82
Conversation
Previously, when the model returned stopReason "error"/"aborted", query() surfaced a system error message and ended the session with no way to retry or resume — Query.steer() was also a no-op stub. Both are now wired into pi-agent-core's Agent instance: - steer(message) queues a steering message via agent.steer(). - continue(message?) resumes a session that has gone idle. With no argument it retries the exact turn the model failed to answer (dropping the failed assistant message so the transcript ends on the user/tool-result message it never responded to, per pi-agent-core's continue() contract). With a message, it queues a follow-up turn. The internal message channel can now reopen after finishing, since a completed query's async iterator may be resumed by continue().
shreyas-lyzr
left a comment
shreyas-lyzr
left a comment
There was a problem hiding this comment.
Five issues found — two confirming the bugs noted in the PR description (one with a more precise failure path), three new ones. Inline comments below.
| throw new Error("Cannot continue: agent not yet initialized (query hasn't started)"); | ||
| } | ||
| if (agent.signal) { | ||
| throw new Error("Cannot continue: agent is still processing a turn"); |
There was a problem hiding this comment.
Bug 1 (confirmed, fix incomplete): this guard fires a false positive when continue() is called from within an agent_end event listener — the most natural point to retry after a failure.
The sequence: runWithLifecycle emits agent_end via processEvents, which awaits each subscriber in order. One subscriber calls channel.finish(), which synchronously resolves the pending pull() promise, scheduling the consumer's for-await continuation as a microtask. That microtask can run before processEvents returns, meaning finishRun() (which clears activeRun and therefore agent.signal) hasn't run yet. agent.signal is still a live, non-aborted AbortSignal — truthy — so this line throws even though the run is effectively done.
Two options:
await agent.waitForIdle()before the guard — but that changes semantics for a genuinely concurrent call (silent wait instead of immediate throw).- Let pi-agent-core's own guard fire (
"Agent is already processing"), catch it, and re-throw with the friendlier message — skippingchannel.reopen()on that path.
| if (message !== undefined) { | ||
| pushMsg({ type: "user", content: message }); | ||
| agent.followUp(toUserMessage(message)); | ||
| } else { |
There was a problem hiding this comment.
Bug 2 (confirmed, triggering condition narrower than stated): agent.followUp(msg) queues the message, but pi-agent-core's agent.continue() only drains the follow-up queue when the last message role in state.messages is "assistant". If the last role is anything else, continue() falls through to runContinuation(), which ignores the queued follow-up entirely — the new message is silently lost and the LLM re-answers the existing context.
In the current code paths this is handled: handleRunFailure unconditionally pushes an error assistant message, so the last role is always "assistant" after any failure. The scenario from the PR description — "401 before any assistant message is recorded" — is not reachable via the current StreamFn contract (errors surface as a final encoded message, not as a thrown exception). The bug is real but latent; it would resurface if a future pi-agent-core version changes that invariant, or if the agent state is externally modified before continue() is called.
| let agent: Agent | undefined; | ||
|
|
||
| // OpenTelemetry session span — opened immediately so it covers agent | ||
| // load + prompt + cleanup. Closed in the IIFE's finally so every exit |
There was a problem hiding this comment.
Bug 3 (new, moderate): every continue() call re-emits a session_start system message into the channel. The agent_start event fires again at the top of runWithLifecycle -> runAgentLoopContinue for both paths (with and without a new message), and the subscriber that handles agent_start unconditionally emits { type: "system", subtype: "session_start" }. A query that fails and is continued twice will produce three session_start and three session_end messages within the same Query lifetime.
Any consumer that treats session_start as a once-per-query lifecycle marker — initialising UI state, starting timers, logging a session — will misfire on every continue() call. Consider gating the emit on a flag (e.g. let sessionStartEmitted = false) so it only fires once per Query.
| try { | ||
| if (message !== undefined) { | ||
| pushMsg({ type: "user", content: message }); | ||
| agent.followUp(toUserMessage(message)); |
There was a problem hiding this comment.
Bug 4 (new, minor but observable): pushMsg({ type: "user", content: message }) is called before agent.followUp() and before await agent.continue(). If the subsequent agent.continue() call rejects (e.g. pi-agent-core throws "Cannot continue from message role: assistant" due to unexpected state), the catch block pushes an error system message and calls channel.finish(). The consumer's messages() now contains the user message as if it was processed by the LLM, when it never was — a ghost entry that persists in collectedMessages for the lifetime of the query.
Consider pushing the user message only after agent.followUp() and agent.continue() resolve successfully, so the consumer-visible transcript stays consistent with what was actually sent.
| } | ||
|
|
||
| await otelContext.with(_session.ctx, () => (agent as Agent).continue()); | ||
| } catch (err: any) { |
There was a problem hiding this comment.
Bug 5 (new, minor): _session.end() is called in the IIFE's finally block, which runs when the first run ends — before any continue() call. Child spans created during a continued run (gen_ai.chat spans, tool spans) are created under _session.ctx via otelContext.with(_session.ctx, ...) here, but by this point the parent session span is already ended. This produces broken telemetry traces with out-of-order or orphaned child spans.
The session span should remain open until the Query itself is done. One approach: move _session.end() into a finalise() helper called from both the normal channel.finish() path and the return()/throw() methods of the generator, guarded by a flag so it fires exactly once.
2 participants
Summary
query()had no way to recover when the model failed to return output. When a turn ended withstopReason: "error"(or"aborted"), the SDK just emitted a system error message and ended the session — no retry, no continuation.Query.steer()was also a no-op stub (steer(_message) {}), despite the underlying@mariozechner/pi-agent-coreAgentalready exposing workingsteer()/followUp()/continue()primitives that were simply never wired up.This PR wires both up:
steer(message)now actually queues a steering message viaagent.steer(...)instead of doing nothing. Effective while the agent is still actively running a turn.continue(message?)(new) resumes a session that has gone idle — most importantly after a failed turn:agent.continue()re-sends it (matchespi-agent-core's documentedcontinue()contract: "Continue from the current transcript. The last message must be a user or tool-result message.").agent.followUp(...)and resumes.reopen()so a query whose async iterator already reachedagent_end/done: truecan resume streaming oncecontinue()is called.Test plan
npm run build(tsc) — clean, no type errorsnpm test— all existing 27 tests pass unchanged; added 4 new tests covering thecontinue()/steer()guard paths (method presence, rejecting before the agent has loaded, matching priorsteer()behavior)pi-agent-core's realAgentclass with a fakestreamFnthat fails once then succeeds — confirmedcontinue()produces a clean transcript (user message + successful assistant reply) rather than leaving the failed message in place