Wolo/context unification#971
Draft
wolo-lab wants to merge 7 commits into
Draft
Conversation
wolo-lab
force-pushed
the
wolo/output-validation-finalize
branch
2 times, most recently
from
bb51d73 to
96a1101
Compare
…urface Foundation for context unification (folding the workflow NodeContext into the unified agent.Context, per go/orcas-rfc-580 and the adk-python model where Context wraps an InvocationContext and carries tool, callback, and workflow-node state). - agent.Context gains Path() and RunID() accessors. They are meaningful for workflow graph nodes and return "" for plain tool/callback contexts. - commonContext (the single concrete Context impl) now forwards the remaining InvocationContext methods (Agent, Memory, Session, RunConfig, EndInvocation, Ended, ResumedInput, WithContext) to its wrapped InvocationContext, so a single unified Context value also satisfies InvocationContext. This lets one value flow wherever an InvocationContext is expected, mirroring adk-python's Context. - Path()/RunID() forward to the wrapped InvocationContext when it carries node information, else return "". - callbackContextWrapper forwards the two new accessors. Purely additive: no signature changes yet. The Node.Run switch to agent.Context comes in a follow-up commit.
Second additive step of context unification. Introduce the agent-level abstractions that let the unified Context carry workflow-node state and schedule dynamic children, without the agent package importing workflow. - NodeScheduler: an opaque interface (implemented by the workflow layer) stored on Context, the Go analog of adk-python's ScheduleDynamicNode protocol carried on Context. - NodeRunOptions: plain data describing a dynamic child run (run id, sub-branch, override branch, use-as-output), so agent can describe the call without importing workflow's functional option type. - agent.Context gains NodeScheduler() alongside the previously added Path()/RunID(). - commonContext gains node-state fields (path, runID, resumeInputs, nodeScheduler); Path/RunID/ResumedInput/NodeScheduler now read them, and WithContext preserves them. - NewNodeContext(ic, path, runID, resumeInputs, sched, actions) builds the unified Context for a workflow node activation (full Context + satisfies InvocationContext), with artifact-delta tracking. - callbackContextWrapper forwards NodeScheduler(). Still additive: no existing callers changed, nothing constructs a node context via the new path yet. The workflow wiring and the Node.Run signature switch follow.
Complete the context unification (go/orcas-rfc-580) for workflow nodes: Node.Run now receives the unified agent.Context instead of a separate workflow.NodeContext, matching adk-python where BaseNode.run takes the unified Context (which wraps an InvocationContext and carries tool, callback, and workflow-node state). agent package: - agent.Context now embeds InvocationContext, so a single unified value exposes the full invocation surface (Agent/Memory/Session/RunConfig/ ...) plus callback, tool, and node accessors. commonContext and callbackContextWrapper forward the InvocationContext methods. - The node surface (Path, RunID, NodeScheduler) and NewNodeContext were added in the preceding commits; this commit makes the workflow layer build and use them. workflow package: - Node.Run signature: agent.InvocationContext -> agent.Context, across the Node interface and every node implementation (Function, Tool, Agent, Join, Dynamic, Workflow, ParallelWorker, Start). - Removed workflow.NodeContext / nodeContext / newNodeContext / newDynamicNodeContext. The per-node context is now built via agent.NewNodeContext. - DynamicFn takes agent.Context. RunNode takes agent.Context and reaches the scheduler via ctx.NodeScheduler() instead of a *nodeContext type assertion. - dynamicSubScheduler implements agent.NodeScheduler (ScheduleNode), adapting agent.NodeRunOptions to the internal runNodeOptions; it keeps the workflow-only state (outputForAncestors, resume inputs) that must not live on the agent package. - The Go-context bridge now carries a *nodeBridge (the node's agent.Context plus workflow-only outputForAncestors/resumeInputs). NodeContextFromGoContext returns agent.Context. The bridge is stashed on static activations, dynamic children, and the dynamic orchestrator body so tools can recover a node context whose NodeScheduler can RunNode. Each bridge is pointed at the final (rewrapped) context. Tests/examples updated for the new signatures; mock contexts gain the InvocationContext methods now required of agent.Context. The context_test assertion that a CallbackContext is not an InvocationContext is inverted to match the unified design.
Adjust the context unification to match adk-python more closely: the unified agent.Context now WRAPS an InvocationContext and is deliberately NOT itself an InvocationContext, instead of embedding it. Rationale: embedding made agent.Context an InvocationContext (IS-A), exposing the full invocation surface (EndInvocation, WithContext, RunConfig, ...) on every callback/tool/node context — broader and more surprising than adk-python, whose Context(ReadonlyContext) holds an InvocationContext as a private field and re-exposes only selected properties. Changes: - agent.Context no longer embeds InvocationContext. It adds InvocationContext() (the accessor, analogous to Python's get_invocation_context()) plus the selected surface node bodies need (Agent, Memory, Session, RunConfig, Ended, ResumedInput) and a WithContext that returns Context (not InvocationContext). - commonContext / callbackContextWrapper expose InvocationContext() and no longer satisfy InvocationContext; WithContext returns Context. - Workflow call sites that need the raw InvocationContext (withBranch, NewNodeContext, NewToolContext, NewRequestInputEvent, startNodeSpan, sub-workflow RunNode, FunctionNode user fn) now call ctx.InvocationContext(). - Tests/examples/mocks updated: reentry/hitl callbacks that read ResumedInput take agent.Context (so resume payloads on the wrapper are visible); mock contexts expose InvocationContext() and a Context-returning WithContext. The context_test assertion confirms a CallbackContext is NOT an InvocationContext (the two are now provably disjoint via differing WithContext return types). Full build and tests pass (the A2A cleanup test is pre-existing flaky under full-suite load; it passes in isolation).
Remove the redundant resumeInputs field from nodeBridge. Resume payloads already ride on the unified agent.Context; expose them via a new ResumeInputs() accessor (mirrors adk-python's Context.resume_inputs) so the dynamic node reads ctx.ResumeInputs() instead of carrying a second copy on the bridge. nodeBridge now carries only: - ctx: the node's agent.Context, recovered by tools running inside an LlmAgent node (NodeContextFromGoContext) to reach the NodeScheduler — the essential reason the bridge exists. - outputForAncestors: workflow-only output-delegation state that cannot live on agent.Context (the agent package must not import workflow). agent.Context gains ResumeInputs() map[string]any (commonContext + callbackContextWrapper). Mock contexts in tests implement it. HITL resume (hitl_simple, dynamic/hitl) verified end-to-end; full build and tests pass.
Tighten visibility of the workflow-only scheduler so it is not part of
the public agent API, mirroring adk-python where _workflow_scheduler is
a private Context field touched only by run_node.
V1 — off the interface:
- Remove NodeScheduler() from the agent.Context interface. Implementers
(including every test mock) no longer need it; the public Context
surface for nodes is now just Path() and RunID() (both public in
adk-python too).
V3a — types out of the public agent package:
- Move the NodeScheduler interface and NodeRunOptions struct from agent
to package workflow (their natural home).
- commonContext now carries the scheduler as an opaque `any` token
(agent never calls it; it only stores and returns it). NewNodeContext
takes `sched any`. The concrete commonContext.NodeScheduler() any
accessor remains exported so package workflow can recover the token,
but it is not on the Context interface.
- workflow.RunNode recovers the token via a local capability interface
(nodeScheduled{ NodeScheduler() any }) and asserts it to
workflow.NodeScheduler.
Net: agent's public API has no scheduling concepts; the scheduler is an
opaque token agent cannot interpret — as close to Python's private
field as Go allows across a package boundary. The 4 test mocks drop
their NodeScheduler() method. Full build, vet, tests, and the dynamic /
HITL samples pass.
Seal the public API: NodeContextFromGoContext was the only exported symbol of the node-context bridge, and it is an internal mechanism (consumed only by single_turn_tool to reach the surrounding node's scheduler). Making it unexported means a future full bridge removal / event-architecture change touches no public symbol. - NodeContextFromGoContext -> nodeContextFromGoContext (unexported). - RunNode now recovers the scheduler internally via schedulerFor(ctx): it first checks whether ctx itself carries a scheduler (a dynamic node body), then falls back to the node context stashed on ctx's embedded Go context (the tool/callback-context-inside-a-node case). The bridge recovery is now an implementation detail of RunNode. - single_turn_tool calls workflow.RunNode(toolCtx, ...) directly instead of recovering the node context itself. Behavior is unchanged (single_turn agent-as-node still propagates child events through the orchestrator). The public workflow API no longer exposes any bridge recovery; only RunNode + the Mode:single_turn configuration remain user-facing. Tests in package workflow use the unexported name; full build and tests pass.
wolo-lab
force-pushed
the
wolo/context-unification
branch
from
5d2c70c to
1190931
Compare
This file contains hidden or bidirectional Unicode text that may be interpreted or compiled differently than what appears below. To review, open the file in an editor that reveals hidden Unicode characters.
Learn more about bidirectional Unicode characters
Sign up for free
to join this conversation on GitHub.
Already have an account?
Sign in to comment
1 participant
Add this suggestion to a batch that can be applied as a single commit.This suggestion is invalid because no changes were made to the code.Suggestions cannot be applied while the pull request is closed.Suggestions cannot be applied while viewing a subset of changes.Only one suggestion per line can be applied in a batch.Add this suggestion to a batch that can be applied as a single commit.Applying suggestions on deleted lines is not supported.You must change the existing code in this line in order to create a valid suggestion.Outdated suggestions cannot be applied.This suggestion has been applied or marked resolved.Suggestions cannot be applied from pending reviews.Suggestions cannot be applied on multi-line comments.Suggestions cannot be applied while the pull request is queued to merge.Suggestion cannot be applied right now. Please check back later.
Please ensure you have read the contribution guide before creating a pull request.
Link to Issue or Description of Change
1. Link to an existing issue (if applicable):
2. Or, if no issue exists, describe the change:
If applicable, please follow the issue templates to provide as much detail as
possible.
Problem:
A clear and concise description of what the problem is.
Solution:
A clear and concise description of what you want to happen and why you choose
this solution.
Testing Plan
Please describe the tests that you ran to verify your changes. This is required
for all PRs that are not small documentation or typo fixes.
Unit Tests:
Please include a summary of passed go test results.
Manual End-to-End (E2E) Tests:
Please provide instructions on how to manually test your changes, including any
necessary setup or configuration. Please provide logs or screenshots to help
reviewers better understand the fix.
Checklist
Additional context
Add any other context or screenshots about the feature request here.