From 68fa067b4c0c831399b5b2574925236b3b3aadb6 Mon Sep 17 00:00:00 2001 From: Antonio Lobato Date: Fri, 21 Aug 2026 16:24:07 -0700 Subject: [PATCH 1/2] feat(workflow): tail-only loop break, final-step outcome report, artifact handoff MIME-Version: 1.0 Content-Type: text/plain; charset=UTF-8 Content-Transfer-Encoding: 8bit Restores three workflow behaviors and does each the ADK-native way. 1. Only the tail step of a loop can break out early. exit_loop is now attached to the last inner step only, not every inner step. It sets Actions.Escalate — the only way an ask tool ends a loopagent — and no other tool touches Actions, so withholding it makes an early break structurally impossible for the non-tail steps. This replaces the old end_turn decision guard, which detected the violation after the fact and re-prompted; now it cannot happen. The non-tail loop instruction is reworded to say the step cannot end the loop. 2. finish_workflow is attached again. The tool and the Progress capture both survived #136, but nothing attached the tool to a step once the graph became one session, so a run reported no artifacts. It is now attached to the final step (WorkflowStepTools). This is how the user learns what a run produced — PR links, tickets — so it is not optional. 3. Artifacts pass structured data between steps. Every step gets save_artifact (native, writes through ctx.Artifacts().Save — ADK ships only load_artifacts) and load_artifacts. This works now, where it did not before #136: the whole graph runs as one runner invocation, so the runner's ArtifactService spans every step. This is the ADK-native replacement for the ask/plans notes directories. Mechanics: CompileWorkflow hands each step a StepRole{InLoop, IsTail, IsFinal} through its ToolsBuilder. IsTail gates exit_loop, IsFinal gates finish_workflow. The TUI builder (cmd/ask/workflow_graph.go) appends tools.WorkflowStepTools directly; the headless builder passes WorkflowStep/WorkflowFinalStep flags to the tool factory, which appends them in BuildCoreTools (pkg/engine can't import pkg/tools). Tests: StepRole is computed correctly for linear/loop/final-loop shapes; WorkflowStepTools attaches save+load always and finish only on the final step; save_artifact writes name+content through ctx.Artifacts().Save and errors (a real Go error, so retryandreflect sees it) with no service or no name; the loop instruction tells only the tail step how to break, and only the final step to call finish_workflow. Co-Authored-By: Claude Opus 4.8 --- CLAUDE.md | 53 ++++++++++++---- cmd/ask/workflow_graph.go | 9 ++- pkg/engine/run.go | 7 +++ pkg/engine/workflow_run.go | 6 +- pkg/tools/artifact.go | 68 +++++++++++++++++++++ pkg/tools/artifact_test.go | 115 +++++++++++++++++++++++++++++++++++ pkg/tools/core.go | 6 +- pkg/workflow/compile.go | 36 +++++++++-- pkg/workflow/compile_test.go | 101 ++++++++++++++++++++++++++++++ pkg/workflow/runner.go | 25 ++++++-- 10 files changed, 397 insertions(+), 29 deletions(-) create mode 100644 pkg/tools/artifact.go create mode 100644 pkg/tools/artifact_test.go diff --git a/CLAUDE.md b/CLAUDE.md index 0119f5c..ed7d126 100644 --- a/CLAUDE.md +++ b/CLAUDE.md @@ -99,7 +99,8 @@ One `package main`, one file per concern. | `workflow_store.go` | Three-scope workflow persistence: user (ask.json) + repo (`/.ask/workflows/*.json`, committed) + global (`~/.config/ask/workflows/*.json`, machine-local, visible from every project); merged global-first listing (personal-wins), ambiguity-strict name resolution, dir sync on save, cross-scope copy. | | `workflows_screen.go` | Workflows builder screen — list/steps/step editor levels with multi-line prompt textarea. `e` on a selected workflow opens that same textarea (workflow-scoped `promptTarget=="description"`) to edit the workflow's free-text Description; it commits to `workflowDef.Description` and shows as the steps-pane subtitle. | | `workflows_picker.go` | Small centred modal popped on `f` (issues) / `Ctrl+F` (chat) to pick which workflow to run. | -| `pkg/workflow/compile.go` | `CompileWorkflow` — a `Def` becomes an ADK graph: `AgentNode` per step, `loopagent`+`exitlooptool` per loop, `IncludeContentsNone` + `InstructionProvider` per step agent, per-node `RetryConfig`. Returns `Compiled` with the agent-name → step-index map the progress adapter needs. | +| `pkg/workflow/compile.go` | `CompileWorkflow` — a `Def` becomes an ADK graph: `AgentNode` per step, `loopagent`+`exitlooptool` (tail step only) per loop, `IncludeContentsNone` + `InstructionProvider` per step agent, per-node `RetryConfig`. Hands each step a `StepRole{InLoop, IsTail, IsFinal}` so the `ToolsBuilder` attaches position-dependent tools. Returns `Compiled` with the agent-name → step-index map the progress adapter needs. | +| `pkg/tools/artifact.go` | `SaveArtifactTool` (writes through `ctx.Artifacts().Save`; ADK ships only `load_artifacts`) and `WorkflowStepTools(env, isFinal)` — save/load on every step for data handoff, `finish_workflow` on the final step for the user-facing outcome report. | | `pkg/workflow/progress.go` | `Progress` — ADK events → `RunnerListener` callbacks, driven only by events that actually arrived. | | `cmd/ask/workflow_graph.go` | TUI runner: one session for the whole graph, agent swapped for the compiled workflow. | | `pkg/engine/workflow_run.go` | Headless runner + `WorkflowGraphAgent` + `IngestWorkflowMemory`. | @@ -438,14 +439,22 @@ were deleted. Don't add a second one. - A top-level agent step becomes an `AgentNode` wrapping an `llmagent`; nodes are chained `Start -> n0 -> n1 -> …`. - A `kind: "loop"` step becomes an `AgentNode` wrapping ADK's - `loopagent`, whose sub-agents are the inner steps, each carrying - `exitlooptool`. A step breaks the loop by calling `exit_loop` (it sets - `Actions.Escalate`, which is what `loopagent` watches for); otherwise - the loop runs to `MaxIterations`. There is no `decision` argument any - more — loop control is ADK's tool, not an `end_turn` field. + `loopagent`, whose sub-agents are the inner steps. **Only the tail + (last) inner step carries `exitlooptool`**, so only it can break the + loop early: `exit_loop` sets `Actions.Escalate` (what `loopagent` + watches), no other ask tool touches `Actions`, so withholding the tool + from the non-tail steps makes an early break structurally impossible + for them — this replaces the old `end_turn` decision guard, which + caught the violation after the fact and re-prompted. Otherwise the loop + runs to `MaxIterations`. There is no `decision` argument — loop control + is ADK's tool, not an `end_turn` field. - `NodeConfig.RetryConfig` gives per-node retry, replacing the runner's hand-rolled `stepErrorRetry` loop. +`CompileWorkflow` hands each step a `StepRole{InLoop, IsTail, IsFinal}` +via its `ToolsBuilder`, which is what decides the position-dependent +tools: `IsTail` gates `exit_loop`, `IsFinal` gates `finish_workflow`. + Two `llmagent` settings carry the semantics and MUST NOT be dropped: - **`IncludeContents: IncludeContentsNone`** is what isolates a step. @@ -500,8 +509,29 @@ Every step should call `end_turn` once with a `summary` (1-3 sentences); it becomes the step's line in the workflow log via `stepSummaryLine`. A step that ends without it is NOT re-prompted any more — the whole remind/re-prompt machinery is gone — its log line falls back to the first -line of its own output. `finish_workflow` still reports the run's -outcome; the runner reads it from `env.PendingFinishData`. +line of its own output. + +`finish_workflow` reports the run's outcome and the artifacts it produced +(PR links, tickets) to the user. It is attached to the **final step +only** (`WorkflowStepTools(env, isFinal)` in pkg/tools/artifact.go); +`Progress` captures the call straight from the event stream +(progress.go), and the TUI also reads `env.PendingFinishData` as a +backup. Nothing attaches it otherwise — an earlier regression left it +attached to no step at all once the graph became one session. + +### Passing data between steps (artifacts) + +A node's text output is the implicit handoff to the next node. For +structured data — a plan, a diff, notes — a step calls the native +`save_artifact` tool; a later step loads it by name with ADK's +`load_artifacts`. Both are attached to every workflow step by +`WorkflowStepTools`. This works because the whole graph runs as ONE +runner invocation, so the runner's `ArtifactService` (set in +`RunnerBuilder`) spans every step — a save in step 1 is visible to a load +in step 3. ADK ships `load_artifacts` but no save tool; `SaveArtifactTool` +(pkg/tools/artifact.go) is the missing half, writing through +`ctx.Artifacts().Save`. This is the ADK-native replacement for the old +`ask/plans/` notes directories. ### Step instruction assembly @@ -524,9 +554,10 @@ node's input, and `IncludeContentsNone` is what lets the step see it. `ask/plans/` is gone, along with `plans.go`, `clear_plans`, and the `RemindFixPlanDir` re-prompt. Step-to-step handoff is the graph's node -output. Durable "what we learned" goes to `pkg/memory` — `RunWorkflow` -calls `engine.IngestWorkflowMemory` on a clean finish, which is what the -notes directories were badly approximating. +output plus ADK artifacts (see above). Durable "what we learned" goes to +`pkg/memory` — `RunWorkflow` calls `engine.IngestWorkflowMemory` on a +clean finish, which is what the notes directories were badly +approximating. ### Builder screen (`Ctrl+W` / `/workflows`) diff --git a/cmd/ask/workflow_graph.go b/cmd/ask/workflow_graph.go index a3e9066..a6443b5 100644 --- a/cmd/ask/workflow_graph.go +++ b/cmd/ask/workflow_graph.go @@ -7,6 +7,7 @@ import ( "github.com/Cidan/ask/pkg/engine" "github.com/Cidan/ask/pkg/providers" + "github.com/Cidan/ask/pkg/tools" "github.com/Cidan/ask/pkg/workflow" adkmodel "google.golang.org/adk/v2/model" adktool "google.golang.org/adk/v2/tool" @@ -127,10 +128,12 @@ func tuiWorkflowCompileConfig(sess *agentSession, def workflow.Def, src workflow ModelBuilder: func(ctx context.Context, step workflow.Step) (adkmodel.LLM, error) { return workflowStepModel(ctx, sess, step) }, - ToolsBuilder: func(ctx context.Context, step workflow.Step, inLoop bool) ([]adktool.Tool, error) { - return engine.AsADKTools(sess.currentTools()) + ToolsBuilder: func(ctx context.Context, step workflow.Step, role workflow.StepRole) ([]adktool.Tool, error) { + list := append([]tools.Tool(nil), sess.currentTools()...) + list = append(list, tools.WorkflowStepTools(sess.env, role.IsFinal)...) + return engine.AsADKTools(list) }, - ToolsetsBuilder: func(ctx context.Context, step workflow.Step, inLoop bool) ([]adktool.Toolset, error) { + ToolsetsBuilder: func(ctx context.Context, step workflow.Step, role workflow.StepRole) ([]adktool.Toolset, error) { var toolsets []adktool.Toolset if sess.mcp != nil { toolsets = append(toolsets, sess.mcp.Toolsets()...) diff --git a/pkg/engine/run.go b/pkg/engine/run.go index 5b66185..0247ed6 100644 --- a/pkg/engine/run.go +++ b/pkg/engine/run.go @@ -105,6 +105,13 @@ type ToolFactoryArgs struct { EventListener EventListener InteractionHandler InteractionHandler AttachWebSearch bool + // WorkflowStep attaches the workflow-step tools (save_artifact, + // load_artifacts) so a step can pass data to a later one. + WorkflowStep bool + // WorkflowFinalStep also attaches finish_workflow, which reports the + // run's outcome and created artifacts to the user. Only meaningful + // with WorkflowStep. + WorkflowFinalStep bool } // ToolFactory builds a slice of Tools for an engine turn. diff --git a/pkg/engine/workflow_run.go b/pkg/engine/workflow_run.go index 5add45d..adc4a2c 100644 --- a/pkg/engine/workflow_run.go +++ b/pkg/engine/workflow_run.go @@ -27,7 +27,7 @@ func WorkflowCompileConfig(e *Engine, cwd string, tabID int, def workflow.Def, s ModelBuilder: func(ctx context.Context, step workflow.Step) (adkmodel.LLM, error) { return buildStepModel(ctx, e, step) }, - ToolsBuilder: func(ctx context.Context, step workflow.Step, inLoop bool) ([]tool.Tool, error) { + ToolsBuilder: func(ctx context.Context, step workflow.Step, role workflow.StepRole) ([]tool.Tool, error) { var agentTools []Tool if tf := GetDefaultToolFactory(); tf != nil { agentTools = tf(ToolFactoryArgs{ @@ -37,11 +37,13 @@ func WorkflowCompileConfig(e *Engine, cwd string, tabID int, def workflow.Def, s EventListener: e.opts.EventListener, InteractionHandler: e.opts.InteractionHandler, AttachWebSearch: true, + WorkflowStep: true, + WorkflowFinalStep: role.IsFinal, }) } return AsADKTools(agentTools) }, - ToolsetsBuilder: func(ctx context.Context, step workflow.Step, inLoop bool) ([]tool.Toolset, error) { + ToolsetsBuilder: func(ctx context.Context, step workflow.Step, role workflow.StepRole) ([]tool.Toolset, error) { var toolsets []tool.Toolset if skillTS, err := NewSkillToolset(ctx, cwd); err == nil && skillTS != nil { toolsets = append(toolsets, skillTS) diff --git a/pkg/tools/artifact.go b/pkg/tools/artifact.go new file mode 100644 index 0000000..f469034 --- /dev/null +++ b/pkg/tools/artifact.go @@ -0,0 +1,68 @@ +package tools + +import ( + "errors" + "strings" + + "google.golang.org/adk/v2/agent" + "google.golang.org/adk/v2/tool/loadartifactstool" + "google.golang.org/genai" +) + +// SaveArtifactParams is the save_artifact tool's input. +type SaveArtifactParams struct { + Name string `json:"name" jsonschema:"artifact name, e.g. plan.md — later steps load it by this name"` + Content string `json:"content" jsonschema:"the artifact's text content"` + Description string `json:"description" jsonschema:"one short human-readable phrase (under 10 words) telling the user what this call is doing"` +} + +// SaveArtifactResult is the save_artifact tool's response. +type SaveArtifactResult struct { + Name string `json:"name,omitempty"` + Version int64 `json:"version,omitempty" jsonschema:"the saved version number"` +} + +// SaveArtifactTool lets a workflow step hand structured data to a later +// step. The artifact lives in the run's ArtifactService, which spans the +// whole graph because the graph runs as one runner invocation, so a step +// saves plan.md and a downstream step reads it back with load_artifacts. +// +// ADK ships load_artifacts but no save tool; this is the missing half. +func SaveArtifactTool() Tool { + return NewTypedTool( + "save_artifact", + "Save a named artifact (a plan, a diff, notes) so a later workflow step can load it by name. Use this to hand structured output forward rather than restating it.", + func(ctx agent.Context, p SaveArtifactParams) (SaveArtifactResult, error) { + name := strings.TrimSpace(p.Name) + if name == "" { + return SaveArtifactResult{}, errors.New("name is required") + } + arts := ctx.Artifacts() + if arts == nil { + return SaveArtifactResult{}, errors.New("artifacts are not available in this context") + } + resp, err := arts.Save(ctx, name, genai.NewPartFromText(p.Content)) + if err != nil { + return SaveArtifactResult{}, err + } + var version int64 + if resp != nil { + version = resp.Version + } + return SaveArtifactResult{Name: name, Version: version}, nil + }, + ) +} + +// WorkflowStepTools returns the position-dependent tools every workflow +// step needs: save_artifact and load_artifacts on every step so steps can +// pass data forward, plus finish_workflow on the final step only, which +// reports the run's outcome and created artifacts (PR links, tickets) to +// the user. +func WorkflowStepTools(env *ToolEnv, isFinal bool) []Tool { + out := []Tool{SaveArtifactTool(), loadartifactstool.New()} + if isFinal { + out = append(out, FinishWorkflowTool(env)) + } + return out +} diff --git a/pkg/tools/artifact_test.go b/pkg/tools/artifact_test.go new file mode 100644 index 0000000..2780a4e --- /dev/null +++ b/pkg/tools/artifact_test.go @@ -0,0 +1,115 @@ +package tools + +import ( + "context" + "strings" + "testing" + + "google.golang.org/adk/v2/agent" + "google.golang.org/adk/v2/artifact" + "google.golang.org/genai" +) + +// A non-final step can pass data forward (save/load) but must not be able +// to report the run's outcome; only the final step gets finish_workflow. +func TestWorkflowStepTools(t *testing.T) { + env := NewToolEnv(t.TempDir(), 1, true, false, nil, nil) + + names := func(ts []Tool) map[string]bool { + m := map[string]bool{} + for _, tl := range ts { + m[tl.Name()] = true + } + return m + } + + mid := names(WorkflowStepTools(env, false)) + if !mid["save_artifact"] || !mid["load_artifacts"] { + t.Errorf("every workflow step needs save/load, got %v", mid) + } + if mid["finish_workflow"] { + t.Error("a non-final step must not get finish_workflow") + } + + final := names(WorkflowStepTools(env, true)) + if !final["save_artifact"] || !final["load_artifacts"] || !final["finish_workflow"] { + t.Errorf("the final step needs save/load/finish, got %v", final) + } +} + +// save_artifact reports a real error (not a silent field) when the +// context has no artifact service, so the failure reaches the model. +func TestSaveArtifactTool_NoArtifactService(t *testing.T) { + _, err := runTypedTool[SaveArtifactResult](t, SaveArtifactTool(), + SaveArtifactParams{Name: "plan.md", Content: "hi", Description: "save"}) + if err == nil || !strings.Contains(err.Error(), "artifacts are not available") { + t.Errorf("save without an artifact service should error, got %v", err) + } +} + +func TestSaveArtifactTool_RequiresName(t *testing.T) { + _, err := runTypedTool[SaveArtifactResult](t, SaveArtifactTool(), + SaveArtifactParams{Name: " ", Content: "hi", Description: "save"}) + if err == nil || !strings.Contains(err.Error(), "name is required") { + t.Errorf("save without a name should error, got %v", err) + } +} + +// fakeArtifacts is a minimal agent.Artifacts backed by a map, so a test +// can watch what save_artifact writes without standing up a runner. +type fakeArtifacts struct{ store map[string]*genai.Part } + +func (f *fakeArtifacts) Save(_ context.Context, name string, data *genai.Part) (*artifact.SaveResponse, error) { + f.store[name] = data + return &artifact.SaveResponse{Version: int64(len(f.store))}, nil +} +func (f *fakeArtifacts) Load(_ context.Context, name string) (*artifact.LoadResponse, error) { + return &artifact.LoadResponse{Part: f.store[name]}, nil +} +func (f *fakeArtifacts) LoadVersion(_ context.Context, name string, _ int) (*artifact.LoadResponse, error) { + return &artifact.LoadResponse{Part: f.store[name]}, nil +} +func (f *fakeArtifacts) List(context.Context) (*artifact.ListResponse, error) { + names := make([]string, 0, len(f.store)) + for n := range f.store { + names = append(names, n) + } + return &artifact.ListResponse{FileNames: names}, nil +} + +// artifactsCtx overrides just Artifacts() on a real agent.Context. +type artifactsCtx struct { + agent.Context + arts agent.Artifacts +} + +func (a artifactsCtx) Artifacts() agent.Artifacts { return a.arts } + +// save_artifact writes the content to the run's artifact store under the +// given name — the handoff a later step reads back with load_artifacts. +func TestSaveArtifactTool_WritesToArtifactStore(t *testing.T) { + arts := &fakeArtifacts{store: map[string]*genai.Part{}} + ctx := artifactsCtx{Context: testAgentCtx(), arts: arts} + + runner, ok := SaveArtifactTool().(interface { + Run(agent.Context, any) (map[string]any, error) + }) + if !ok { + t.Fatal("save_artifact does not implement the ADK Run contract") + } + res, err := runner.Run(ctx, map[string]any{ + "name": "plan.md", + "content": "the plan body", + "description": "save the plan", + }) + if err != nil { + t.Fatalf("save_artifact: %v", err) + } + if res["name"] != "plan.md" { + t.Errorf("result name = %v, want plan.md", res["name"]) + } + saved := arts.store["plan.md"] + if saved == nil || saved.Text != "the plan body" { + t.Errorf("artifact store did not receive the content: %+v", saved) + } +} diff --git a/pkg/tools/core.go b/pkg/tools/core.go index a06c968..9210435 100644 --- a/pkg/tools/core.go +++ b/pkg/tools/core.go @@ -28,7 +28,11 @@ func BuildCoreTools(args engine.ToolFactoryArgs, attachWebSearch bool) []Tool { registryFunc := func() []Tool { return nil } - return CoreTools(env, registryFunc, attachWebSearch) + core := CoreTools(env, registryFunc, attachWebSearch) + if args.WorkflowStep { + core = append(core, WorkflowStepTools(env, args.WorkflowFinalStep)...) + } + return core } // BuildSubagentTools constructs the centralized toolset configured for an isolated subagent session. diff --git a/pkg/workflow/compile.go b/pkg/workflow/compile.go index 561f423..73554da 100644 --- a/pkg/workflow/compile.go +++ b/pkg/workflow/compile.go @@ -29,8 +29,10 @@ type WorkflowAgentConfig struct { // ModelBuilder resolves the LLM for one step. Required. ModelBuilder func(ctx context.Context, step Step) (model.LLM, error) // ToolsBuilder and ToolsetsBuilder supply the step's tool surface. - ToolsBuilder func(ctx context.Context, step Step, inLoop bool) ([]tool.Tool, error) - ToolsetsBuilder func(ctx context.Context, step Step, inLoop bool) ([]tool.Toolset, error) + // StepRole tells them where the step sits, so the builder can attach + // finish_workflow to the final step only, and so on. + ToolsBuilder func(ctx context.Context, step Step, role StepRole) ([]tool.Tool, error) + ToolsetsBuilder func(ctx context.Context, step Step, role StepRole) ([]tool.Toolset, error) // InstructionBuilder renders the step's system instruction. Defaults // to BuildStepInstruction. InstructionBuilder func(step Step, src Source, pc *StepPromptCtx) string @@ -44,6 +46,20 @@ type WorkflowAgentConfig struct { // RetryConfig, which also handles the backoff. const workflowDefaultMaxRetries = 3 +// StepRole tells a ToolsBuilder where a step sits in the workflow, so it +// can decide which position-dependent tools to attach. +type StepRole struct { + // InLoop is true when the step runs inside a loop container. + InLoop bool + // IsTail is true when the step is the last inner step of its loop — + // the only step allowed to break the loop early via exit_loop. + IsTail bool + // IsFinal is true when the step is the last thing the whole workflow + // runs (the last top-level step, or the tail of a final loop). Only + // this step gets finish_workflow. + IsFinal bool +} + // Compiled is a Def rendered as an executable ADK graph, plus the lookup // the event adapter needs to attribute events back to steps. type Compiled struct { @@ -186,16 +202,24 @@ func buildStepAgent(ctx context.Context, cfg WorkflowAgentConfig, step Step, ste return nil, fmt.Errorf("workflow compile: model for step %q: %w", step.Name, err) } - inLoop := loop != nil + role := StepRole{InLoop: loop != nil, IsFinal: isFinal} + if loop != nil { + role.IsTail = loop.IsTail + } var tools []tool.Tool if cfg.ToolsBuilder != nil { - tools, err = cfg.ToolsBuilder(ctx, step, inLoop) + tools, err = cfg.ToolsBuilder(ctx, step, role) if err != nil { return nil, fmt.Errorf("workflow compile: tools for step %q: %w", step.Name, err) } } - if inLoop { + // Only the tail step of a loop may break out early. exit_loop sets + // Actions.Escalate, which is the only way an ask tool ends a + // loopagent, and no other tool touches Actions — so withholding it + // from the non-tail steps makes an early break structurally + // impossible for them, rather than something caught after the fact. + if role.InLoop && role.IsTail { tools, err = withExitLoopTool(tools) if err != nil { return nil, fmt.Errorf("workflow compile: step %q: %w", step.Name, err) @@ -204,7 +228,7 @@ func buildStepAgent(ctx context.Context, cfg WorkflowAgentConfig, step Step, ste var toolsets []tool.Toolset if cfg.ToolsetsBuilder != nil { - toolsets, err = cfg.ToolsetsBuilder(ctx, step, inLoop) + toolsets, err = cfg.ToolsetsBuilder(ctx, step, role) if err != nil { return nil, fmt.Errorf("workflow compile: toolsets for step %q: %w", step.Name, err) } diff --git a/pkg/workflow/compile_test.go b/pkg/workflow/compile_test.go index 7ad1b0f..cbe1a39 100644 --- a/pkg/workflow/compile_test.go +++ b/pkg/workflow/compile_test.go @@ -231,3 +231,104 @@ func TestBuildStepInstruction_LoopFraming(t *testing.T) { t.Errorf("loop control moved to exit_loop; instruction should not mention a decision arg: %q", got) } } + +// Only the tail step is told how to break, because only the tail step has +// the exit_loop tool. A non-tail step is told it cannot end the loop. +func TestBuildStepInstruction_NonTailLoopStepCannotBreak(t *testing.T) { + src := NewTextSource(1, "src") + step := Step{Name: "edit", Prompt: "Edit."} + loop := &LoopPromptCtx{Name: "fix", MaxIterations: 5, IsTail: false} + + got := BuildStepInstruction(step, src, &StepPromptCtx{Loop: loop}) + if strings.Contains(got, "exit_loop") { + t.Errorf("a non-tail step must not be told about exit_loop: %q", got) + } + if !strings.Contains(got, "NOT its last step") { + t.Errorf("a non-tail step should be told it cannot end the loop: %q", got) + } +} + +// The final step is told to report the run's outcome with finish_workflow. +func TestBuildStepInstruction_FinalStepReportsWithFinishWorkflow(t *testing.T) { + src := NewTextSource(1, "src") + step := Step{Name: "ship", Prompt: "Open the PR."} + + final := BuildStepInstruction(step, src, &StepPromptCtx{IsWorkflowFinalStep: true}) + if !strings.Contains(final, "finish_workflow") { + t.Errorf("the final step must be told to call finish_workflow: %q", final) + } + + mid := BuildStepInstruction(step, src, &StepPromptCtx{}) + if strings.Contains(mid, "finish_workflow") { + t.Errorf("a non-final step must not be told about finish_workflow: %q", mid) + } + // Every step is told it can pass data forward. + if !strings.Contains(mid, "save_artifact") { + t.Errorf("every step should learn about save_artifact: %q", mid) + } +} + +// recordRoles compiles def with a ToolsBuilder that records the StepRole +// handed to each step, keyed by step name. +func recordRoles(t *testing.T, def Def) map[string]StepRole { + t.Helper() + roles := map[string]StepRole{} + cfg := testCompileConfig(def) + cfg.ToolsBuilder = func(_ context.Context, step Step, role StepRole) ([]tool.Tool, error) { + roles[step.Name] = role + return nil, nil + } + if _, err := CompileWorkflow(context.Background(), cfg); err != nil { + t.Fatalf("compile: %v", err) + } + return roles +} + +// The role handed to each step is what decides its position-dependent +// tools: IsTail gates exit_loop (only the tail step may break a loop), +// IsFinal gates finish_workflow (only the last step reports the outcome). +func TestCompileWorkflow_StepRoles(t *testing.T) { + def := Def{Name: "wf", Steps: []Step{ + {Name: "a", Prompt: "a"}, + {Name: "fix", Kind: "loop", Steps: []Step{ + {Name: "b", Prompt: "b"}, + {Name: "c", Prompt: "c"}, + {Name: "d", Prompt: "d"}, + }}, + {Name: "e", Prompt: "e"}, + }} + roles := recordRoles(t, def) + + want := map[string]StepRole{ + "a": {InLoop: false, IsTail: false, IsFinal: false}, + "b": {InLoop: true, IsTail: false, IsFinal: false}, + "c": {InLoop: true, IsTail: false, IsFinal: false}, + "d": {InLoop: true, IsTail: true, IsFinal: false}, // tail, but the loop is not the final step + "e": {InLoop: false, IsTail: false, IsFinal: true}, + } + for name, w := range want { + if roles[name] != w { + t.Errorf("role[%q] = %+v, want %+v", name, roles[name], w) + } + } +} + +// When the workflow ends in a loop, the loop's tail step is BOTH the +// break point and the final step, so it gets exit_loop and finish_workflow. +func TestCompileWorkflow_FinalLoopTailIsFinal(t *testing.T) { + def := Def{Name: "wf", Steps: []Step{ + {Name: "a", Prompt: "a"}, + {Name: "fix", Kind: "loop", Steps: []Step{ + {Name: "b", Prompt: "b"}, + {Name: "c", Prompt: "c"}, + }}, + }} + roles := recordRoles(t, def) + + if r := roles["b"]; r.IsTail || r.IsFinal { + t.Errorf("non-tail loop step b = %+v, want neither tail nor final", r) + } + if r := roles["c"]; !r.IsTail || !r.IsFinal { + t.Errorf("tail of the final loop c = %+v, want both tail and final", r) + } +} diff --git a/pkg/workflow/runner.go b/pkg/workflow/runner.go index a86bfad..95e9899 100644 --- a/pkg/workflow/runner.go +++ b/pkg/workflow/runner.go @@ -76,8 +76,17 @@ func BuildStepInstruction(step Step, source Source, pc *StepPromptCtx) string { b.WriteString(ref) } var loop *LoopPromptCtx + isFinal := false if pc != nil { loop = pc.Loop + isFinal = pc.IsWorkflowFinalStep + } + b.WriteString("\n\nTo hand structured output to a later step — a plan, a diff, notes — call save_artifact " + + "with a name, and the later step loads it with load_artifacts. Prefer this over restating large output in your summary.") + if isFinal { + b.WriteString(" You are the FINAL step of this workflow: before you finish, call finish_workflow with a " + + "description of the outcome and the list of artifacts it produced — every PR, issue, or link the user " + + "needs. This is how the user learns what the run created.") } b.WriteString("\n\n") b.WriteString(EndTurnInstructionBlock(loop)) @@ -103,12 +112,16 @@ func EndTurnInstructionBlock(loop *LoopPromptCtx) string { "a `summary` of 1-3 sentences describing what you did and the outcome. This records your progress in the " + "workflow log; it does not cut your turn short.") if loop != nil { - b.WriteString(" You are running inside a loop: when the loop's exit goal above is met, call the " + - "exit_loop tool to end the loop. If it is not met, do not call exit_loop — the loop advances to its " + - "next iteration on its own, and stops by itself after the iteration limit.") - if !loop.IsTail { - b.WriteString(" Later steps in this iteration still have work to do, so only call exit_loop if the " + - "goal is already fully met.") + if loop.IsTail { + // Only the tail step is given the exit_loop tool, so only the + // tail step is told how to break. + b.WriteString(" You are the last step of this loop iteration: when the loop's exit goal above is met, " + + "call the exit_loop tool to end the loop. If it is not met, do not call exit_loop — the loop advances " + + "to its next iteration on its own, and stops by itself after the iteration limit.") + } else { + b.WriteString(" You are running inside a loop but you are NOT its last step, so you cannot end the loop — " + + "only the final step of the iteration decides whether to continue or stop. Do your part and end your " + + "turn.") } } return b.String() From f7925de85ff2c9c0c367aa8c9fbf4c5adabdb8ca Mon Sep 17 00:00:00 2001 From: Antonio Lobato Date: Fri, 21 Aug 2026 16:29:58 -0700 Subject: [PATCH 2/2] test(workflow): end-to-end artifact handoff through the real ADK graph MIME-Version: 1.0 Content-Type: text/plain; charset=UTF-8 Content-Transfer-Encoding: 8bit A two-step workflow on the real ADK runner: step 1 saves an artifact, step 2 loads it. Proves the whole seam the earlier tests only covered in pieces — the tool factory attaches save_artifact/load_artifacts to every step, the graph runs as one runner invocation so the ArtifactService spans both steps, and the saved content reaches step 2's model. Lives in the external engine_test package so it can import pkg/tools (whose init registers the tool factory) without the tools -> engine import cycle. A scripted model.LLM drives one step at a time by the marker in each step's system instruction. The assertion is airtight against a false pass: the secret must be ABSENT from the reader's step handoff (the saver's node output is 'saved the plan', and IncludeContentsNone keeps its tool calls out of the reader's context) and PRESENT only after load_artifacts resolves. A negative control with a mismatched artifact name fails as expected. Co-Authored-By: Claude Opus 4.8 --- .../workflow_artifact_integration_test.go | 177 ++++++++++++++++++ 1 file changed, 177 insertions(+) create mode 100644 pkg/engine/workflow_artifact_integration_test.go diff --git a/pkg/engine/workflow_artifact_integration_test.go b/pkg/engine/workflow_artifact_integration_test.go new file mode 100644 index 0000000..49c8dcb --- /dev/null +++ b/pkg/engine/workflow_artifact_integration_test.go @@ -0,0 +1,177 @@ +package engine_test + +import ( + "context" + "iter" + "strings" + "sync" + "testing" + + "github.com/Cidan/ask/pkg/config" + "github.com/Cidan/ask/pkg/engine" + "github.com/Cidan/ask/pkg/providers" + // Imported for its init(): registers the tool factory, so workflow + // steps get the real core toolset plus save_artifact / load_artifacts. + _ "github.com/Cidan/ask/pkg/tools" + "github.com/Cidan/ask/pkg/workflow" + adkmodel "google.golang.org/adk/v2/model" + "google.golang.org/genai" +) + +// scriptedModel drives one step at a time by inspecting the marker in the +// step's system instruction and counting calls per marker. It is a full +// model.LLM, not the internal test mock, because this test lives in the +// external engine_test package to avoid the tools -> engine import cycle. +type scriptedModel struct { + mu sync.Mutex + calls map[string]int + // readerHandoff is what the reader saw on its FIRST call — the step + // handoff, before any artifact was loaded. + readerHandoff string + // readerContent captures every request text the reader step saw on + // its second call, so the test can assert the artifact crossed over. + readerContent []string +} + +func (m *scriptedModel) Name() string { return "scripted" } + +func (m *scriptedModel) GenerateContent(_ context.Context, req *adkmodel.LLMRequest, _ bool) iter.Seq2[*adkmodel.LLMResponse, error] { + m.mu.Lock() + defer m.mu.Unlock() + + sys := systemText(req) + switch { + case strings.Contains(sys, "SAVE_STEP_MARKER"): + n := m.calls["save"] + m.calls["save"]++ + if n == 0 { + return one(functionCall("save_artifact", map[string]any{ + "name": "plan.md", + "content": "SECRET_PLAN_BODY", + "description": "save the plan", + })) + } + return one(finalText("saved the plan")) + + case strings.Contains(sys, "READ_STEP_MARKER"): + n := m.calls["read"] + m.calls["read"]++ + if n == 0 { + m.readerHandoff = requestText(req) + return one(functionCall("load_artifacts", map[string]any{ + "artifact_names": []any{"plan.md"}, + })) + } + // Second call: load_artifacts' ProcessRequest has injected the + // artifact content into the request by now. + m.readerContent = append(m.readerContent, requestText(req)) + return one(finalText("read the plan")) + } + return one(finalText("done")) +} + +func systemText(req *adkmodel.LLMRequest) string { + if req == nil || req.Config == nil || req.Config.SystemInstruction == nil { + return "" + } + var b strings.Builder + for _, p := range req.Config.SystemInstruction.Parts { + if p != nil { + b.WriteString(p.Text) + } + } + return b.String() +} + +func requestText(req *adkmodel.LLMRequest) string { + var b strings.Builder + for _, c := range req.Contents { + if c == nil { + continue + } + for _, p := range c.Parts { + if p != nil && p.Text != "" { + b.WriteString(p.Text) + b.WriteByte('\n') + } + } + } + return b.String() +} + +func functionCall(name string, args map[string]any) *adkmodel.LLMResponse { + return &adkmodel.LLMResponse{ + Content: &genai.Content{ + Role: genai.RoleModel, + Parts: []*genai.Part{genai.NewPartFromFunctionCall(name, args)}, + }, + } +} + +func finalText(text string) *adkmodel.LLMResponse { + return &adkmodel.LLMResponse{ + Content: &genai.Content{Role: genai.RoleModel, Parts: []*genai.Part{genai.NewPartFromText(text)}}, + FinishReason: genai.FinishReasonStop, + } +} + +func one(resp *adkmodel.LLMResponse) iter.Seq2[*adkmodel.LLMResponse, error] { + return func(yield func(*adkmodel.LLMResponse, error) bool) { yield(resp, nil) } +} + +// TestWorkflowArtifactHandoff is the full round trip: a two-step workflow +// on the real ADK graph, where step 1 saves an artifact and step 2 loads +// it. It proves the wiring end to end — the tool factory attaches +// save_artifact/load_artifacts to every step, the graph runs as one +// runner invocation so the ArtifactService spans both steps, and the +// content saved by step 1 reaches step 2's model. +func TestWorkflowArtifactHandoff(t *testing.T) { + t.Setenv("HOME", t.TempDir()) + tmp := t.TempDir() + + scripted := &scriptedModel{calls: map[string]int{}} + origBuilder := engine.ModelBuilder + engine.ModelBuilder = func(_ context.Context, _ *providers.AgentProviderSpec, _ config.Config, _ string) (adkmodel.LLM, error) { + return scripted, nil + } + defer func() { engine.ModelBuilder = origBuilder }() + + eng := engine.New(engine.Options{ + Config: config.Config{Provider: "vertex"}, + InteractionHandler: engine.HeadlessInteractionHandler{AutoApproveTools: true}, + }) + + def := workflow.Def{ + Name: "artifact-handoff", + Steps: []workflow.Step{ + {Name: "saver", Prompt: "SAVE_STEP_MARKER: save the plan as an artifact."}, + {Name: "reader", Prompt: "READ_STEP_MARKER: load the plan artifact and use it."}, + }, + } + src := workflow.NewTextSource(1, "Artifact Handoff Source") + + if err := eng.RunWorkflow(context.Background(), tmp, 1, def, src); err != nil { + t.Fatalf("RunWorkflow: %v", err) + } + + scripted.mu.Lock() + defer scripted.mu.Unlock() + + if scripted.calls["save"] == 0 { + t.Fatal("the saver step never ran") + } + if len(scripted.readerContent) == 0 { + t.Fatal("the reader step never reached its second call, so load_artifacts never resolved") + } + // The secret must NOT arrive through the normal step handoff — the + // saver's node output is "saved the plan", and IncludeContentsNone + // keeps the saver's tool calls out of the reader's context. So its + // only path to the reader is load_artifacts. + if strings.Contains(scripted.readerHandoff, "SECRET_PLAN_BODY") { + t.Fatalf("the secret leaked through the step handoff, not the artifact:\n%s", scripted.readerHandoff) + } + joined := strings.Join(scripted.readerContent, "\n") + if !strings.Contains(joined, "SECRET_PLAN_BODY") { + t.Errorf("the artifact saved by step 1 did not reach step 2's model via load_artifacts:\n%s", joined) + } +}