diff --git a/.github/actions/setup/action.yml b/.github/actions/setup/action.yml index efc6b5cd6..df4994cec 100644 --- a/.github/actions/setup/action.yml +++ b/.github/actions/setup/action.yml @@ -11,13 +11,13 @@ runs: steps: - name: Set up Go id: setup-go - uses: actions/setup-go@d35c59abb061a4a6fb18e82ac0862c26744d6ab5 # v5.5.0 + uses: actions/setup-go@4a3601121dd01d1626a1e23e37211e3254c1c06c # v6.4.0 with: go-version-file: 'go.mod' cache: false # wrapper for actions/cache that doesn't support all functionality - name: Load Go cache - uses: actions/cache@5a3ec84eff668545956fd18022155c47e93e2684 # v4.2.3 + uses: actions/cache@27d5ce7f107fe9357f9df03efb73ab90386fccae # v5.0.5 with: path: | ~/.cache/go-build diff --git a/.github/dependabot.yml b/.github/dependabot.yml new file mode 100644 index 000000000..a0544629c --- /dev/null +++ b/.github/dependabot.yml @@ -0,0 +1,18 @@ +version: 2 +updates: + - package-ecosystem: "gomod" + directory: "/" + schedule: + interval: "weekly" + day: "monday" + time: "02:00" + commit-message: + prefix: "chore" + include: "scope" + labels: + - "dependencies" + reviewers: + - "wolo-lab" + allow: + - dependency-name: "google.golang.org/genai" + open-pull-requests-limit: 10 diff --git a/.github/workflows/go.yml b/.github/workflows/go.yml index fedc37dc9..121284133 100644 --- a/.github/workflows/go.yml +++ b/.github/workflows/go.yml @@ -7,10 +7,10 @@ name: Go on: push: - branches: [ "main" ] + branches: [ "main", "v2" ] pull_request: - branches: [ "main" ] + branches: [ "main", "v2" ] workflow_dispatch: diff --git a/.github/workflows/nightly.yml b/.github/workflows/nightly.yml index 9222fb13c..78c641520 100644 --- a/.github/workflows/nightly.yml +++ b/.github/workflows/nightly.yml @@ -22,7 +22,11 @@ jobs: runs-on: ubuntu-latest steps: - name: Run govulncheck - uses: golang/govulncheck-action@b625fbe08f3bccbe446d94fbf87fcc875a4f50ee # v1.0.4 + # NOTE: pinned to master HEAD instead of a tag. The latest release (v1.0.4) + # internally uses actions/checkout@v4.1.1 and actions/setup-go@v5.0.0 (Node 20), + # which are deprecated. master uses checkout v6.0.2 + setup-go v6.2.0 (Node 24). + # TODO: re-pin to a tagged release once upstream tags one past v1.0.4. + uses: golang/govulncheck-action@31f7c5463448f83528bd771c2d978d940080c9fd # master, post-v1.0.4 (Node 24) with: go-version-file: go.mod repo-checkout: true \ No newline at end of file diff --git a/.golangci.yml b/.golangci.yml index dd478fe62..4c597cef0 100644 --- a/.golangci.yml +++ b/.golangci.yml @@ -53,6 +53,7 @@ linters: # Changes to the default config: - -QF1001 # Exclude "Apply De Morgan's law" - -QF1008 # Exclude "Omit embedded fields from selector expression" + - -SA1019 # Exclude to keep the old a2a compatibility, can remove for ADK v2 exclusions: rules: diff --git a/AGENTS.md b/AGENTS.md new file mode 100644 index 000000000..b75cab3a8 --- /dev/null +++ b/AGENTS.md @@ -0,0 +1,149 @@ +# AGENTS.md + +Context for AI coding agents (Claude Code, Gemini CLI, Cursor, Copilot, etc.) +working in the ADK Go repository. Human contributors should start with +CONTRIBUTING.md. + +## Project overview + +ADK Go (`google.golang.org/adk`) is an open-source, code-first Go toolkit for +building, evaluating, and deploying AI agents. It is model-agnostic but +optimized for Gemini, and is one of three ADK implementations — Go, Python, and +Java — that share a conceptual model but are independent codebases. Requires +Go 1.25+. + +## Setup & core commands + +Run from the repo root. These match what CI enforces (CI also passes `-v`): + +- Build: `go build -mod=readonly ./...` +- Test: `go test -race -mod=readonly -count=1 -shuffle=on ./...` +- Single pkg: `go test -race ./agent/...` +- Lint: `golangci-lint run` (golangci-lint v2; CI pins v2.3.1; config in `.golangci.yml`) +- Tidy check: `go mod tidy -diff` (must print nothing) +- Format: `golangci-lint fmt` (applies gofumpt + goimports per config) + +## Definition of done + +A change is complete only when all of these pass locally: + +1. `go build` (above) succeeds. +2. `go test` (above) is green. +3. `golangci-lint run` reports no findings. +4. `go mod tidy -diff` prints nothing. +5. New/changed behavior has tests; a bug fix has a test that reproduces the bug. +6. Every new Go file starts with the Apache 2.0 license header (enforced by `goheader`). + +## Repository layout + +- `agent/` Agent interface + types (`llmagent`, `workflowagents`, `remoteagent`) +- `runner/` Execution engine that drives the run loop +- `model/` LLM abstraction (`gemini`, `apigee`) +- `tool/` Tool/Toolset interface + built-in tools (incl. `skilltoolset/`, `mcptoolset/`) +- `session/` Conversation state + events +- `memory/`, `artifact/` Long-term memory and file/data services +- `plugin/` Cross-cutting lifecycle hooks +- `server/` HTTP servers (`adkrest` is primary; `adka2a`, `agentengine`) +- `cmd/` CLI (`adkgo`) and server launchers +- `telemetry/`, `util/` Public helper packages +- `internal/` Private packages — NOT public API; `internal/httprr` is vendored +- `examples/` Runnable example agents (quickstart, tools, a2a, skills, …) + +## Conventions & idioms + +- **Streaming:** agent runs return `iter.Seq2[*session.Event, error]`; consume + with `for event, err := range … {}`. Don't collect events into a slice. +- **Interface-first:** public packages expose interfaces (`Agent`, `Tool`, + `Toolset`, `Service`); concrete impls live in sub-packages or `internal/`. +- **Callbacks over subclassing** (`Before*`/`After*` for Agent/Model/Tool); + returning non-nil from a `Before` callback short-circuits execution. +- **Errors:** wrap with `fmt.Errorf("…: %w", err)`. Tool confirmation uses + sentinel errors (e.g. `tool.ErrConfirmationRequired`). +- Prefer an existing helper over a new one; keep packages small and focused. + +## Minimal example + +```go +model, err := gemini.NewModel(ctx, "gemini-2.5-flash", + &genai.ClientConfig{APIKey: os.Getenv("GOOGLE_API_KEY")}) +// handle err +a, err := llmagent.New(llmagent.Config{ + Name: "assistant", + Model: model, + Instruction: "You are a helpful assistant.", + Tools: []tool.Tool{ /* ... */ }, +}) +// handle err +r, err := runner.New(runner.Config{ + AppName: "my-app", + Agent: a, + SessionService: session.InMemoryService(), + AutoCreateSession: true, +}) +// handle err +msg := genai.NewContentFromText("Hello", genai.RoleUser) +for event, err := range r.Run(ctx, userID, sessionID, msg, agent.RunConfig{}) { + // handle err; read event.LLMResponse.Content +} +``` + +See `examples/quickstart` for a full runnable program. + +## Extending the framework + +- **Add a tool:** wrap a Go function with + `functiontool.New[Args, Results](cfg, handler)` (Args/Results are structs), or + implement the `tool.Tool` interface for full control. +- **Add a toolset:** implement `tool.Toolset`; its `Tools(ctx)` may return + different tools per invocation. +- **Add an agent type:** follow the `agent/workflowagents/*` packages; construct + agents via `llmagent.New` / `agent.New`, not by implementing `agent.Agent` + directly. +- **Add cross-cutting behavior:** register a `plugin.New(plugin.Config{...})` + hook (`Before*`/`After*` for run/agent/model/tool) instead of editing the loop. + +## Testing + +- Tests run **offline by default**: LLM HTTP traffic is replayed from + `testdata/*.httprr` via `internal/httprr`. Never add live model or network + calls to tests. +- To (re)record a package's traffic, supply real credentials (e.g. + `GOOGLE_API_KEY`) and run `go generate .//...` (it runs + `go test -httprecord=…`); commit the updated `testdata/*.httprr`. +- Prefer table-driven tests; shared helpers live in `internal/testutil`. + +## Boundaries + +**Always** +- Run build, tests, lint, and `go mod tidy -diff` before declaring done. +- Keep PRs small and focused — one concern per PR. +- Add or update tests for the code you change. + +**Ask first** +- Adding or upgrading a dependency (`go.mod`). +- Changing a high-fan-in package (`session`, `agent`, `model`, `tool`, + `runner`) — prefer additive, backward-compatible changes. +- Any change to the public API surface, and any breaking change. + +**Never** +- Break the public API — keep changes backward-compatible. +- Edit vendored code (`internal/httprr`) or commit secrets / API keys. +- Add tests that make live LLM or network calls. + +## PRs & commits + +See `CONTRIBUTING.md` for the full process and CLA. Key points for agents: +most PRs (beyond trivial docs/typos) need a linked issue; include a **Testing +Plan**; attach logs or screenshots for behavior changes (Runner output / ADK Web). + +## Alignment with adk-python + +[adk-python](https://github.com/google/adk-python) is the source of truth for +feature behavior. When porting or validating a feature, check parity with the +Python implementation. + +## Resources + +- Docs: https://google.github.io/adk-docs/ +- Examples: `./examples` +- Java ADK: https://github.com/google/adk-java diff --git a/CLAUDE.md b/CLAUDE.md new file mode 100644 index 000000000..543a9c15f --- /dev/null +++ b/CLAUDE.md @@ -0,0 +1 @@ +See [AGENTS.md](./AGENTS.md) for project context, commands, and contribution guidelines for AI coding agents. diff --git a/GEMINI.md b/GEMINI.md new file mode 100644 index 000000000..543a9c15f --- /dev/null +++ b/GEMINI.md @@ -0,0 +1 @@ +See [AGENTS.md](./AGENTS.md) for project context, commands, and contribution guidelines for AI coding agents. diff --git a/agent/agent.go b/agent/agent.go index e94ed4f5d..f34893805 100644 --- a/agent/agent.go +++ b/agent/agent.go @@ -249,11 +249,8 @@ func runBeforeAgentCallbacks(ctx InvocationContext) (*session.Event, error) { agent := ctx.Agent() pluginManager := pluginManagerFromContext(ctx) - callbackCtx := &callbackContext{ - Context: ctx, - invocationContext: ctx, - actions: &session.EventActions{StateDelta: make(map[string]any), ArtifactDelta: make(map[string]int64)}, - } + actions := &session.EventActions{StateDelta: make(map[string]any), ArtifactDelta: make(map[string]int64)} + callbackCtx := NewCallbackContext(ctx, actions) if pluginManager != nil { content, err := pluginManager.RunBeforeAgentCallback(callbackCtx) @@ -261,13 +258,13 @@ func runBeforeAgentCallbacks(ctx InvocationContext) (*session.Event, error) { return nil, fmt.Errorf("failed to run plugin before agent callback: %w", err) } if content != nil { - event := session.NewEvent(ctx.InvocationID()) + event := session.NewEventWithContext(ctx, ctx.InvocationID()) event.LLMResponse = model.LLMResponse{ Content: content, } event.Author = agent.Name() event.Branch = ctx.Branch() - event.Actions = *callbackCtx.actions + event.Actions = *actions ctx.EndInvocation() return event, nil } @@ -282,23 +279,23 @@ func runBeforeAgentCallbacks(ctx InvocationContext) (*session.Event, error) { continue } - event := session.NewEvent(ctx.InvocationID()) + event := session.NewEventWithContext(ctx, ctx.InvocationID()) event.LLMResponse = model.LLMResponse{ Content: content, } event.Author = agent.Name() event.Branch = ctx.Branch() - event.Actions = *callbackCtx.actions + event.Actions = *actions ctx.EndInvocation() return event, nil } // check if has delta create event with it - if len(callbackCtx.actions.StateDelta) > 0 { - event := session.NewEvent(ctx.InvocationID()) + if len(actions.StateDelta) > 0 { + event := session.NewEventWithContext(ctx, ctx.InvocationID()) event.Author = agent.Name() event.Branch = ctx.Branch() - event.Actions = *callbackCtx.actions + event.Actions = *actions return event, nil } @@ -311,11 +308,8 @@ func runAfterAgentCallbacks(ctx InvocationContext) (*session.Event, error) { agent := ctx.Agent() pluginManager := pluginManagerFromContext(ctx) - callbackCtx := &callbackContext{ - Context: ctx, - invocationContext: ctx, - actions: &session.EventActions{StateDelta: make(map[string]any), ArtifactDelta: make(map[string]int64)}, - } + actions := &session.EventActions{StateDelta: make(map[string]any), ArtifactDelta: make(map[string]int64)} + callbackCtx := NewCallbackContext(ctx, actions) if pluginManager != nil { content, err := pluginManager.RunAfterAgentCallback(callbackCtx) @@ -323,13 +317,13 @@ func runAfterAgentCallbacks(ctx InvocationContext) (*session.Event, error) { return nil, fmt.Errorf("failed to run plugin after agent callback: %w", err) } if content != nil { - event := session.NewEvent(ctx.InvocationID()) + event := session.NewEventWithContext(ctx, ctx.InvocationID()) event.LLMResponse = model.LLMResponse{ Content: content, } event.Author = agent.Name() event.Branch = ctx.Branch() - event.Actions = *callbackCtx.actions + event.Actions = *actions return event, nil } } @@ -343,107 +337,29 @@ func runAfterAgentCallbacks(ctx InvocationContext) (*session.Event, error) { continue } - event := session.NewEvent(ctx.InvocationID()) + event := session.NewEventWithContext(ctx, ctx.InvocationID()) event.LLMResponse = model.LLMResponse{ Content: newContent, } event.Author = agent.Name() event.Branch = ctx.Branch() - event.Actions = *callbackCtx.actions + event.Actions = *actions // TODO set context invocation ended // ctx.invocationEnded = true return event, nil } // check if has delta create event with it - if len(callbackCtx.actions.StateDelta) > 0 { - event := session.NewEvent(ctx.InvocationID()) + if len(actions.StateDelta) > 0 { + event := session.NewEventWithContext(ctx, ctx.InvocationID()) event.Author = agent.Name() event.Branch = ctx.Branch() - event.Actions = *callbackCtx.actions + event.Actions = *actions return event, nil } return nil, nil } -// TODO: unify with internal/context.callbackContext - -type callbackContext struct { - context.Context - invocationContext InvocationContext - actions *session.EventActions -} - -func (c *callbackContext) AgentName() string { - return c.invocationContext.Agent().Name() -} - -func (c *callbackContext) ReadonlyState() session.ReadonlyState { - return c.invocationContext.Session().State() -} - -func (c *callbackContext) State() session.State { - return &callbackContextState{ctx: c} -} - -func (c *callbackContext) Artifacts() Artifacts { - return c.invocationContext.Artifacts() -} - -func (c *callbackContext) InvocationID() string { - return c.invocationContext.InvocationID() -} - -func (c *callbackContext) UserContent() *genai.Content { - return c.invocationContext.UserContent() -} - -// AppName implements CallbackContext. -func (c *callbackContext) AppName() string { - return c.invocationContext.Session().AppName() -} - -// Branch implements CallbackContext. -func (c *callbackContext) Branch() string { - return c.invocationContext.Branch() -} - -// SessionID implements CallbackContext. -func (c *callbackContext) SessionID() string { - return c.invocationContext.Session().ID() -} - -// UserID implements CallbackContext. -func (c *callbackContext) UserID() string { - return c.invocationContext.Session().UserID() -} - -var _ CallbackContext = (*callbackContext)(nil) - -type callbackContextState struct { - ctx *callbackContext -} - -func (c *callbackContextState) Get(key string) (any, error) { - if c.ctx.actions != nil && c.ctx.actions.StateDelta != nil { - if val, ok := c.ctx.actions.StateDelta[key]; ok { - return val, nil - } - } - return c.ctx.invocationContext.Session().State().Get(key) -} - -func (c *callbackContextState) Set(key string, val any) error { - if c.ctx.actions != nil && c.ctx.actions.StateDelta != nil { - c.ctx.actions.StateDelta[key] = val - } - return c.ctx.invocationContext.Session().State().Set(key, val) -} - -func (c *callbackContextState) All() iter.Seq2[string, any] { - return c.ctx.invocationContext.Session().State().All() -} - type invocationContext struct { context.Context diff --git a/agent/callback_context.go b/agent/callback_context.go new file mode 100644 index 000000000..df923247a --- /dev/null +++ b/agent/callback_context.go @@ -0,0 +1,261 @@ +// Copyright 2026 Google LLC +// +// Licensed under the Apache License, Version 2.0 (the "License"); +// you may not use this file except in compliance with the License. +// You may obtain a copy of the License at +// +// http://www.apache.org/licenses/LICENSE-2.0 +// +// Unless required by applicable law or agreed to in writing, software +// distributed under the License is distributed on an "AS IS" BASIS, +// WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. +// See the License for the specific language governing permissions and +// limitations under the License. + +package agent + +import ( + "context" + "fmt" + "iter" + + "google.golang.org/genai" + + "google.golang.org/adk/artifact" + "google.golang.org/adk/memory" + "google.golang.org/adk/platform" + "google.golang.org/adk/session" + "google.golang.org/adk/tool/toolconfirmation" +) + +// NewCallbackContext returns CallbackContext initialized with provided actions. +// actions may be nil; if so, a new session.EventActions is created with empty StateDelta and ArtifactDelta +func NewCallbackContext(ic InvocationContext, actions *session.EventActions) CallbackContext { + actions = prepareEventActions(actions) + cc := &callbackContext{ + Context: ic, + invocationContext: ic, + actions: actions, + artifacts: ic.Artifacts(), + } + return cc +} + +// NewCallbackContextWithArtifactTracking returns CallbackContext initialized with provided actions. +// the returned context's Artifacts().Save(...) wrapper records each saved artifact's version into the underlying +// EventActions.ArtifactDelta so the resulting Event reflects the saves. +// actions may be nil; if so, a new session.EventActions is created with empty StateDelta and ArtifactDelta +func NewCallbackContextWithArtifactTracking(ic InvocationContext, actions *session.EventActions) CallbackContext { + actions = prepareEventActions(actions) + cc := &callbackContext{ + Context: ic, + invocationContext: ic, + actions: actions, + artifacts: &trackedArtifacts{Artifacts: ic.Artifacts(), actions: actions}, + } + return cc +} + +// NewToolContext constructs a ToolContext for a tool execution. +// +// If functionCallID is empty a new UUID is generated. If actions is nil a +// fresh session.EventActions with empty StateDelta and ArtifactDelta is +// allocated; missing sub-maps are populated. The returned ToolContext is +// backed by the same *callbackContext implementation used for CallbackContext, +// so all callback-context semantics (state delta tracking, artifact delta +// tracking, etc.) apply, plus the tool-specific extensions on ToolContext. +func NewToolContext(ic InvocationContext, functionCallID string, actions *session.EventActions, confirmation *toolconfirmation.ToolConfirmation) ToolContext { + if functionCallID == "" { + functionCallID = platform.NewUUID(ic) + } + actions = prepareEventActions(actions) + return &callbackContext{ + Context: ic, + invocationContext: ic, + actions: actions, + artifacts: &trackedArtifacts{Artifacts: ic.Artifacts(), actions: actions}, + functionCallID: functionCallID, + toolConfirmation: confirmation, + } +} + +func prepareEventActions(actions *session.EventActions) *session.EventActions { + if actions == nil { + return &session.EventActions{StateDelta: make(map[string]any), ArtifactDelta: make(map[string]int64)} + } + // create missing maps if needed + if actions.StateDelta == nil { + actions.StateDelta = make(map[string]any) + } + if actions.ArtifactDelta == nil { + actions.ArtifactDelta = make(map[string]int64) + } + return actions +} + +// callbackContext is the single concrete implementation of CallbackContext +// (and, when constructed via NewToolContext, of ToolContext as well). The +// tool-specific methods (FunctionCallID, Actions, SearchMemory, +// ToolConfirmation, RequestConfirmation) are always present on the concrete +// type; they are only meaningful when the context is used as a ToolContext. +type callbackContext struct { + context.Context + invocationContext InvocationContext + artifacts Artifacts + actions *session.EventActions + + // Fields below are only populated by NewToolContext. + functionCallID string + toolConfirmation *toolconfirmation.ToolConfirmation +} + +func (c *callbackContext) AgentName() string { + return c.invocationContext.Agent().Name() +} + +func (c *callbackContext) ReadonlyState() session.ReadonlyState { + return c.invocationContext.Session().State() +} + +func (c *callbackContext) State() session.State { + return &callbackContextState{ctx: c} +} + +func (c *callbackContext) Artifacts() Artifacts { + return c.artifacts +} + +func (c *callbackContext) InvocationID() string { + return c.invocationContext.InvocationID() +} + +func (c *callbackContext) UserContent() *genai.Content { + return c.invocationContext.UserContent() +} + +func (c *callbackContext) AppName() string { + return c.invocationContext.Session().AppName() +} + +func (c *callbackContext) Branch() string { + return c.invocationContext.Branch() +} + +func (c *callbackContext) SessionID() string { + return c.invocationContext.Session().ID() +} + +func (c *callbackContext) UserID() string { + return c.invocationContext.Session().UserID() +} + +var ( + _ CallbackContext = (*callbackContext)(nil) + _ ToolContext = (*callbackContext)(nil) +) + +// --- ToolContext extensions ---------------------------------------------- +// +// The methods below are always present on *callbackContext but only +// meaningful when the context was constructed via NewToolContext (i.e. +// when functionCallID is set). + +// FunctionCallID returns the function call identifier associated with the +// current tool execution, or "" if this context was not constructed for a +// tool call. +func (c *callbackContext) FunctionCallID() string { + return c.functionCallID +} + +// Actions returns the EventActions for the current event. Tools can mutate +// the returned value to influence the agent loop (e.g. state deltas, agent +// transfers). +func (c *callbackContext) Actions() *session.EventActions { + return c.actions +} + +// SearchMemory performs a semantic search on the agent's memory. +func (c *callbackContext) SearchMemory(ctx context.Context, query string) (*memory.SearchResponse, error) { + if c.invocationContext.Memory() == nil { + return nil, fmt.Errorf("memory service is not set") + } + return c.invocationContext.Memory().SearchMemory(ctx, query) +} + +// ToolConfirmation returns the Human-in-the-Loop confirmation handle for the +// current tool execution, or nil if no confirmation is currently associated +// with the call. +func (c *callbackContext) ToolConfirmation() *toolconfirmation.ToolConfirmation { + return c.toolConfirmation +} + +// RequestConfirmation initiates the Human-in-the-Loop (HITL) approval flow +// for the current tool call. It records a pending confirmation in the +// underlying EventActions and sets SkipSummarization so the agent loop halts +// until the user responds. +func (c *callbackContext) RequestConfirmation(hint string, payload any) error { + if c.functionCallID == "" { + return fmt.Errorf("error function call id not set when requesting confirmation for tool") + } + if c.actions.RequestedToolConfirmations == nil { + c.actions.RequestedToolConfirmations = make(map[string]toolconfirmation.ToolConfirmation) + } + c.actions.RequestedToolConfirmations[c.functionCallID] = toolconfirmation.ToolConfirmation{ + Hint: hint, + Confirmed: false, + Payload: payload, + } + // SkipSummarization stops the agent loop after this tool call. Without it, + // the function response event becomes lastEvent and IsFinalResponse() returns + // false (hasFunctionResponses == true), causing the loop to continue. + c.actions.SkipSummarization = true + return nil +} + +// callbackContextState is a session.State implementation backed by the +// callback context's EventActions.StateDelta and the underlying session state. +type callbackContextState struct { + ctx *callbackContext +} + +func (c *callbackContextState) Get(key string) (any, error) { + if c.ctx.actions != nil && c.ctx.actions.StateDelta != nil { + if val, ok := c.ctx.actions.StateDelta[key]; ok { + return val, nil + } + } + return c.ctx.invocationContext.Session().State().Get(key) +} + +func (c *callbackContextState) Set(key string, val any) error { + if c.ctx.actions != nil && c.ctx.actions.StateDelta != nil { + c.ctx.actions.StateDelta[key] = val + } + return c.ctx.invocationContext.Session().State().Set(key, val) +} + +func (c *callbackContextState) All() iter.Seq2[string, any] { + return c.ctx.invocationContext.Session().State().All() +} + +// trackedArtifacts wraps an Artifacts to record each successful Save into the +// supplied EventActions.ArtifactDelta. +type trackedArtifacts struct { + Artifacts + actions *session.EventActions +} + +func (a *trackedArtifacts) Save(ctx context.Context, name string, data *genai.Part) (*artifact.SaveResponse, error) { + resp, err := a.Artifacts.Save(ctx, name, data) + if err != nil { + return resp, err + } + if a.actions != nil { + if a.actions.ArtifactDelta == nil { + a.actions.ArtifactDelta = make(map[string]int64) + } + // TODO: RWLock, check the version stored is newer in case multiple tools save the same file. + a.actions.ArtifactDelta[name] = resp.Version + } + return resp, nil +} diff --git a/agent/context.go b/agent/context.go index 25ea13640..bb40d73ee 100644 --- a/agent/context.go +++ b/agent/context.go @@ -19,7 +19,9 @@ import ( "google.golang.org/genai" + "google.golang.org/adk/memory" "google.golang.org/adk/session" + "google.golang.org/adk/tool/toolconfirmation" ) /* @@ -136,3 +138,63 @@ type CallbackContext interface { Artifacts() Artifacts State() session.State } + +// ToolContext is the context passed to a tool when it is called. It extends +// CallbackContext with tool-specific facilities: access to the originating +// function call, mutable event actions, long-term memory search, and the +// Human-in-the-Loop (HITL) confirmation flow. +type ToolContext interface { + CallbackContext + + // FunctionCallID returns the unique identifier of the function call + // that triggered this tool execution. + FunctionCallID() string + + // Actions returns the EventActions for the current event. This can be + // used by the tool to modify the agent's state, transfer to another + // agent, or perform other actions. + Actions() *session.EventActions + + // SearchMemory performs a semantic search on the agent's memory. + SearchMemory(ctx context.Context, query string) (*memory.SearchResponse, error) + + // ToolConfirmation returns a handler for checking the Human-in-the-Loop + // confirmation status for the current tool context. This should be used + // within a tool's logic *before* performing any sensitive operations that + // require user approval. + // + // Example Usage: + // if confirmation := ctx.ToolConfirmation(); confirmation == nil { + // // Confirmation required, create confirmation or handle appropriately + // ctx.RequestConfirmation("hint", payload) + // } + // + // The returned *toolconfirmation.ToolConfirmation object provides methods + // to check the actual confirmation state. + ToolConfirmation() *toolconfirmation.ToolConfirmation + + // RequestConfirmation initiates the Human-in-the-Loop (HITL) process to + // ask the user for approval before the tool proceeds with a specific + // action. Call this method when a tool needs explicit user consent. + // + // This will typically result in the ADK emitting a special event + // (e.g., a FunctionCall like "adk_request_confirmation") to the client + // application/UI, prompting the user for a decision. + // + // Args: + // - hint: A human-readable string explaining why confirmation is needed. + // This is usually displayed to the user in the confirmation prompt. + // - payload: Any additional data or context about the action requiring + // confirmation. + // + // Returns: + // - nil: If the confirmation request was successfully enqueued or + // initiated within the ADK. This indicates that the process of asking + // the user has begun. It does NOT mean the action is approved. The + // tool's execution will likely pause or be suspended until the user + // responds. + // - error: If there was a failure in initiating the confirmation process + // itself (e.g., invalid arguments, issue with the event system). The + // request to ask the user has not been sent. + RequestConfirmation(hint string, payload any) error +} diff --git a/agent/context_mock.go b/agent/context_mock.go new file mode 100644 index 000000000..9facf40dc --- /dev/null +++ b/agent/context_mock.go @@ -0,0 +1,129 @@ +// Copyright 2026 Google LLC +// +// Licensed under the Apache License, Version 2.0 (the "License"); +// you may not use this file except in compliance with the License. +// You may obtain a copy of the License at +// +// http://www.apache.org/licenses/LICENSE-2.0 +// +// Unless required by applicable law or agreed to in writing, software +// distributed under the License is distributed on an "AS IS" BASIS, +// WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. +// See the License for the specific language governing permissions and +// limitations under the License. + +package agent + +import ( + "context" + "time" + + "google.golang.org/genai" + + "google.golang.org/adk/memory" + "google.golang.org/adk/session" + "google.golang.org/adk/tool/toolconfirmation" +) + +// StrictContextMock is a strict test double for the context interfaces +// ([ToolContext], [CallbackContext], [ReadonlyContext]). +// +// Embed it in a test fake and override only the methods your test actually +// uses. Because it implements the whole surface, embedders keep compiling as +// the interfaces grow. +// +// An un-overridden method panics with "not implemented" — an unexpected call +// fails the test loudly instead of silently returning a zero value. +// +// The exception is the standard library's context.Context methods (Deadline, +// Done, Err and Value): those read from the supplied Ctx rather than panicking, +// so the mock carries a usable context payload. If Ctx is nil they panic like +// everything else. +type StrictContextMock struct { + // Ctx supplies the values returned by Deadline, Done, Err and Value. + Ctx context.Context +} + +func (m *StrictContextMock) ctx() context.Context { + if m.Ctx == nil { + panic("agent.StrictContextMock: Ctx is nil") + } + return m.Ctx +} + +// context.Context methods, served from Ctx instead of panicking. + +// Deadline implements [ToolContext]. +func (m *StrictContextMock) Deadline() (deadline time.Time, ok bool) { return m.ctx().Deadline() } + +// Done implements [ToolContext]. +func (m *StrictContextMock) Done() <-chan struct{} { return m.ctx().Done() } + +// Err implements [ToolContext]. +func (m *StrictContextMock) Err() error { return m.ctx().Err() } + +// Value implements [ToolContext]. +func (m *StrictContextMock) Value(key any) any { return m.ctx().Value(key) } + +// ReadonlyContext methods. + +// UserContent implements [ToolContext]. +func (m *StrictContextMock) UserContent() *genai.Content { panic("not implemented") } + +// InvocationID implements [ToolContext]. +func (m *StrictContextMock) InvocationID() string { panic("not implemented") } + +// AgentName implements [ToolContext]. +func (m *StrictContextMock) AgentName() string { panic("not implemented") } + +// ReadonlyState implements [ToolContext]. +func (m *StrictContextMock) ReadonlyState() session.ReadonlyState { panic("not implemented") } + +// UserID implements [ToolContext]. +func (m *StrictContextMock) UserID() string { panic("not implemented") } + +// AppName implements [ToolContext]. +func (m *StrictContextMock) AppName() string { panic("not implemented") } + +// SessionID implements [ToolContext]. +func (m *StrictContextMock) SessionID() string { panic("not implemented") } + +// Branch implements [ToolContext]. +func (m *StrictContextMock) Branch() string { panic("not implemented") } + +// CallbackContext methods. + +// Artifacts implements [ToolContext]. +func (m *StrictContextMock) Artifacts() Artifacts { panic("not implemented") } + +// State implements [ToolContext]. +func (m *StrictContextMock) State() session.State { panic("not implemented") } + +// ToolContext methods. + +// FunctionCallID implements [ToolContext]. +func (m *StrictContextMock) FunctionCallID() string { panic("not implemented") } + +// Actions implements [ToolContext]. +func (m *StrictContextMock) Actions() *session.EventActions { panic("not implemented") } + +// SearchMemory implements [ToolContext]. +func (m *StrictContextMock) SearchMemory(context.Context, string) (*memory.SearchResponse, error) { + panic("not implemented") +} + +// ToolConfirmation implements [ToolContext]. +func (m *StrictContextMock) ToolConfirmation() *toolconfirmation.ToolConfirmation { + panic("not implemented") +} + +// RequestConfirmation implements [ToolContext]. +func (m *StrictContextMock) RequestConfirmation(hint string, payload any) error { + panic("not implemented") +} + +var ( + _ ToolContext = (*StrictContextMock)(nil) + _ CallbackContext = (*StrictContextMock)(nil) + _ ReadonlyContext = (*StrictContextMock)(nil) +) diff --git a/agent/context_mock_test.go b/agent/context_mock_test.go new file mode 100644 index 000000000..f6346260c --- /dev/null +++ b/agent/context_mock_test.go @@ -0,0 +1,62 @@ +// Copyright 2026 Google LLC +// +// Licensed under the Apache License, Version 2.0 (the "License"); +// you may not use this file except in compliance with the License. +// You may obtain a copy of the License at +// +// http://www.apache.org/licenses/LICENSE-2.0 +// +// Unless required by applicable law or agreed to in writing, software +// distributed under the License is distributed on an "AS IS" BASIS, +// WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. +// See the License for the specific language governing permissions and +// limitations under the License. + +package agent + +import ( + "context" + "testing" +) + +// fakeToolContext shows the intended usage: embed StrictContextMock and +// override only the methods the test needs. This mirrors how out-of-tree +// consumers build a ToolContext test double without tracking every method. +type fakeToolContext struct { + StrictContextMock +} + +var _ ToolContext = (*fakeToolContext)(nil) + +func TestStrictContextMock_ValueDelegatesToCtx(t *testing.T) { + type key struct{} + ctx := context.WithValue(context.Background(), key{}, "v") + + f := &fakeToolContext{StrictContextMock{Ctx: ctx}} + + if got := f.Value(key{}); got != "v" { + t.Errorf("Value() = %v, want %q", got, "v") + } +} + +func TestStrictContextMock_ADKMethodPanics(t *testing.T) { + f := &fakeToolContext{StrictContextMock{Ctx: context.Background()}} + + defer func() { + if r := recover(); r == nil { + t.Error("FunctionCallID() did not panic, want panic for unimplemented method") + } + }() + _ = f.FunctionCallID() +} + +func TestStrictContextMock_NilCtxPanics(t *testing.T) { + f := &fakeToolContext{} + + defer func() { + if r := recover(); r == nil { + t.Error("Value() did not panic with nil Ctx") + } + }() + _ = f.Value("k") +} diff --git a/agent/live.go b/agent/live.go new file mode 100644 index 000000000..cdeaa2c86 --- /dev/null +++ b/agent/live.go @@ -0,0 +1,49 @@ +// Copyright 2026 Google LLC +// +// Licensed under the Apache License, Version 2.0 (the "License"); +// you may not use this file except in compliance with the License. +// You may obtain a copy of the License at +// +// http://www.apache.org/licenses/LICENSE-2.0 +// +// Unless required by applicable law or agreed to in writing, software +// distributed under the License is distributed on an "AS IS" BASIS, +// WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. +// See the License for the specific language governing permissions and +// limitations under the License. + +package agent + +import ( + "google.golang.org/genai" +) + +// LiveSession manages the bidirectional stream for a live session. +type LiveSession interface { + Send(req LiveRequest) error + Close() error +} + +// LiveRequest represents an incoming client event for a live session. +type LiveRequest struct { + // RealtimeInput can be *genai.Blob, *genai.ActivityStart, or *genai.ActivityEnd. + RealtimeInput any + + // Content represents standard text or multimodal content from the user. + // Can also represent a reply to a tool call if it contains a FunctionResponse part. + Content *genai.Content +} + +// LiveRunConfig contains options for configuring a live session. +type LiveRunConfig struct { + ResponseModalities []genai.Modality + SpeechConfig *genai.SpeechConfig + InputAudioTranscription *genai.AudioTranscriptionConfig + OutputAudioTranscription *genai.AudioTranscriptionConfig + RealtimeInputConfig *genai.RealtimeInputConfig + EnableAffectiveDialog bool + Proactivity *genai.ProactivityConfig + SessionResumption *genai.SessionResumptionConfig + SaveLiveBlob bool + MaxLLMCalls int +} diff --git a/agent/llmagent/dynamic_events_test.go b/agent/llmagent/dynamic_events_test.go new file mode 100644 index 000000000..76f9c6fbb --- /dev/null +++ b/agent/llmagent/dynamic_events_test.go @@ -0,0 +1,87 @@ +// Copyright 2025 Google LLC +// +// Licensed under the Apache License, Version 2.0 (the "License"); +// you may not use this file except in compliance with the License. +// You may obtain a copy of the License at +// +// http://www.apache.org/licenses/LICENSE-2.0 +// +// Unless required by applicable law or agreed to in writing, software +// distributed under the License is distributed on an "AS IS" BASIS, +// WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. +// See the License for the specific language governing permissions and +// limitations under the License. + +package llmagent_test + +import ( + "iter" + "slices" + "testing" + + "google.golang.org/genai" + + "google.golang.org/adk/agent" + "google.golang.org/adk/internal/testutil" + "google.golang.org/adk/session" +) + +func TestSessionEvents_YieldedPresence(t *testing.T) { + // Create a custom agent that yields an event and then checks the session events. + customAgent, err := agent.New(agent.Config{ + Name: "test_agent", + Run: func(ctx agent.InvocationContext) iter.Seq2[*session.Event, error] { + return func(yield func(*session.Event, error) bool) { + // 1. Yield a test event + testEvent := session.NewEventWithContext(ctx, ctx.InvocationID()) + testEvent.Content = genai.NewContentFromText("Initial test event", genai.RoleModel) + if !yield(testEvent, nil) { + return + } + + // 2. Check if the event is in the session + var found bool + if ctx.Session() != nil { + for e := range ctx.Session().Events().All() { + if e.Content != nil && len(e.Content.Parts) > 0 { + if e.Content.Parts[0].Text == "Initial test event" { + found = true + break + } + } + } + } + + // 3. Yield the result of the check + resultEvent := session.NewEventWithContext(ctx, ctx.InvocationID()) + if found { + resultEvent.Content = genai.NewContentFromText("Found initial event in session", genai.RoleModel) + } else { + resultEvent.Content = genai.NewContentFromText("Did NOT find initial event in session", genai.RoleModel) + } + yield(resultEvent, nil) + } + }, + }) + if err != nil { + t.Fatalf("failed to create custom agent: %v", err) + } + + runner := testutil.NewTestAgentRunner(t, customAgent) + + // Run the agent + var results []string + for ev, err := range runner.Run(t, "test_session", "Hello") { + if err != nil { + t.Fatalf("run failed: %v", err) + } + if ev.Content != nil && len(ev.Content.Parts) > 0 { + results = append(results, ev.Content.Parts[0].Text) + } + } + + // Verify the result + if !slices.Contains(results, "Found initial event in session") { + t.Errorf("Expected to find initial event in session, but results were: %v", results) + } +} diff --git a/agent/llmagent/llmagent.go b/agent/llmagent/llmagent.go index c3b34493a..ff1149d5a 100644 --- a/agent/llmagent/llmagent.go +++ b/agent/llmagent/llmagent.go @@ -312,7 +312,7 @@ type OnModelErrorCallback func(ctx agent.CallbackContext, llmRequest *model.LLMR // // To modify tool arguments and still run the tool, // update args in place and return (nil, nil). -type BeforeToolCallback func(ctx tool.Context, tool tool.Tool, args map[string]any) (map[string]any, error) +type BeforeToolCallback func(ctx agent.ToolContext, tool tool.Tool, args map[string]any) (map[string]any, error) // AfterToolCallback is a function type executed after a tool's Run method has completed, // regardless of whether the tool returned a result or an error. @@ -321,13 +321,13 @@ type BeforeToolCallback func(ctx tool.Context, tool tool.Tool, args map[string]a // If a callback returns a non-nil result or an error: // - execution of remaining callbacks stops // - the returned result and/or error is used as the final tool output -type AfterToolCallback func(ctx tool.Context, tool tool.Tool, args, result map[string]any, err error) (map[string]any, error) +type AfterToolCallback func(ctx agent.ToolContext, tool tool.Tool, args, result map[string]any, err error) (map[string]any, error) // OnToolErrorCallback that is called when receiving an error response from tool execution. // // If it returns non-nil LLMResponse or error, the actual model response/error // is replaced with the returned response/error. -type OnToolErrorCallback func(ctx tool.Context, tool tool.Tool, args map[string]any, err error) (map[string]any, error) +type OnToolErrorCallback func(ctx agent.ToolContext, tool tool.Tool, args map[string]any, err error) (map[string]any, error) // IncludeContents controls what parts of prior conversation history is received by llmagent. type IncludeContents string @@ -599,6 +599,49 @@ func buildLiveTools(tools []tool.Tool) []*genai.Tool { return []*genai.Tool{{FunctionDeclarations: decls}} } +func (a *llmAgent) RunLive(ctx agent.InvocationContext) (agent.LiveSession, iter.Seq2[*session.Event, error], error) { + ctx = icontext.NewInvocationContext(ctx, icontext.InvocationContextParams{ + Artifacts: ctx.Artifacts(), + Memory: ctx.Memory(), + Session: ctx.Session(), + Branch: ctx.Branch(), + Agent: a, + UserContent: ctx.UserContent(), + RunConfig: ctx.RunConfig(), + InvocationID: ctx.InvocationID(), + }) + + f := &llminternal.Flow{ + Model: a.model, + RequestProcessors: llminternal.DefaultRequestProcessors, + ResponseProcessors: llminternal.DefaultResponseProcessors, + BeforeModelCallbacks: a.beforeModelCallbacks, + AfterModelCallbacks: a.afterModelCallbacks, + OnModelErrorCallbacks: a.onModelErrorCallbacks, + BeforeToolCallbacks: a.beforeToolCallbacks, + AfterToolCallbacks: a.afterToolCallbacks, + OnToolErrorCallbacks: a.onToolErrorCallbacks, + } + + sess, innerIter, err := f.RunLive(ctx) + if err != nil { + return nil, nil, err + } + + wrappedIter := func(yield func(*session.Event, error) bool) { + for ev, err := range innerIter { + if err == nil { + a.maybeSaveOutputToState(ev) + } + if !yield(ev, err) { + return + } + } + } + + return sess, wrappedIter, nil +} + // maybeSaveOutputToState saves the model output to state if needed. skip if the event // was authored by some other agent (e.g. current agent transferred to another agent) func (a *llmAgent) maybeSaveOutputToState(event *session.Event) { diff --git a/agent/llmagent/llmagent_test.go b/agent/llmagent/llmagent_test.go index 1d434ad19..a82a4930f 100644 --- a/agent/llmagent/llmagent_test.go +++ b/agent/llmagent/llmagent_test.go @@ -449,7 +449,7 @@ func TestToolCallback(t *testing.T) { Number int `json:"number"` } - handler := func(_ tool.Context, input Args) (Result, error) { + handler := func(_ agent.ToolContext, input Args) (Result, error) { return Result{Number: 1}, nil } rand, _ := functiontool.New(functiontool.Config{ @@ -468,10 +468,10 @@ func TestToolCallback(t *testing.T) { DisallowTransferToPeers: true, Tools: []tool.Tool{rand}, BeforeToolCallbacks: []llmagent.BeforeToolCallback{ - func(ctx tool.Context, tool tool.Tool, args map[string]any) (map[string]any, error) { + func(ctx agent.ToolContext, tool tool.Tool, args map[string]any) (map[string]any, error) { return nil, nil }, - func(ctx tool.Context, tool tool.Tool, args map[string]any) (map[string]any, error) { + func(ctx agent.ToolContext, tool tool.Tool, args map[string]any) (map[string]any, error) { return map[string]any{"number": "7"}, nil }, }, @@ -504,10 +504,10 @@ func TestToolCallback(t *testing.T) { Tools: []tool.Tool{rand}, BeforeToolCallbacks: []llmagent.BeforeToolCallback{ // Since it retursn non nil, the next callback won't be executed. - func(ctx tool.Context, tool tool.Tool, args map[string]any) (map[string]any, error) { + func(ctx agent.ToolContext, tool tool.Tool, args map[string]any) (map[string]any, error) { return map[string]any{"number": "3"}, nil }, - func(ctx tool.Context, tool tool.Tool, args map[string]any) (map[string]any, error) { + func(ctx agent.ToolContext, tool tool.Tool, args map[string]any) (map[string]any, error) { return map[string]any{"number": "7"}, nil }, }, @@ -539,10 +539,10 @@ func TestToolCallback(t *testing.T) { DisallowTransferToPeers: true, Tools: []tool.Tool{rand}, AfterToolCallbacks: []llmagent.AfterToolCallback{ - func(ctx tool.Context, tool tool.Tool, args, result map[string]any, err error) (map[string]any, error) { + func(ctx agent.ToolContext, tool tool.Tool, args, result map[string]any, err error) (map[string]any, error) { return nil, nil }, - func(ctx tool.Context, tool tool.Tool, args, result map[string]any, err error) (map[string]any, error) { + func(ctx agent.ToolContext, tool tool.Tool, args, result map[string]any, err error) (map[string]any, error) { return map[string]any{"number": "7"}, nil }, }, @@ -575,10 +575,10 @@ func TestToolCallback(t *testing.T) { Tools: []tool.Tool{rand}, AfterToolCallbacks: []llmagent.AfterToolCallback{ // Since it retursn non nil, the next callback won't be executed. - func(ctx tool.Context, tool tool.Tool, args, result map[string]any, err error) (map[string]any, error) { + func(ctx agent.ToolContext, tool tool.Tool, args, result map[string]any, err error) (map[string]any, error) { return map[string]any{"number": "3"}, nil }, - func(ctx tool.Context, tool tool.Tool, args, result map[string]any, err error) (map[string]any, error) { + func(ctx agent.ToolContext, tool tool.Tool, args, result map[string]any, err error) (map[string]any, error) { return map[string]any{"number": "7"}, nil }, }, @@ -610,12 +610,12 @@ func TestToolCallback(t *testing.T) { DisallowTransferToPeers: true, Tools: []tool.Tool{rand}, BeforeToolCallbacks: []llmagent.BeforeToolCallback{ - func(ctx tool.Context, tool tool.Tool, args map[string]any) (map[string]any, error) { + func(ctx agent.ToolContext, tool tool.Tool, args map[string]any) (map[string]any, error) { return map[string]any{"number": "3"}, nil }, }, AfterToolCallbacks: []llmagent.AfterToolCallback{ - func(ctx tool.Context, tool tool.Tool, args, result map[string]any, err error) (map[string]any, error) { + func(ctx agent.ToolContext, tool tool.Tool, args, result map[string]any, err error) (map[string]any, error) { return map[string]any{"number": "7"}, nil }, }, @@ -647,12 +647,12 @@ func TestToolCallback(t *testing.T) { DisallowTransferToPeers: true, Tools: []tool.Tool{rand}, BeforeToolCallbacks: []llmagent.BeforeToolCallback{ - func(ctx tool.Context, tool tool.Tool, args map[string]any) (map[string]any, error) { + func(ctx agent.ToolContext, tool tool.Tool, args map[string]any) (map[string]any, error) { return nil, nil }, }, AfterToolCallbacks: []llmagent.AfterToolCallback{ - func(ctx tool.Context, tool tool.Tool, args, result map[string]any, err error) (map[string]any, error) { + func(ctx agent.ToolContext, tool tool.Tool, args, result map[string]any, err error) (map[string]any, error) { return nil, nil }, }, @@ -857,7 +857,7 @@ func TestFunctionTool(t *testing.T) { } prompt := "what is the sum of 1 + 2?" - handler := func(_ tool.Context, input Args) (Result, error) { + handler := func(_ agent.ToolContext, input Args) (Result, error) { if input.A != 1 || input.B != 2 { t.Errorf("handler received %+v, want {a: 1, b: 2}", input) } diff --git a/agent/llmagent/state_agent_test.go b/agent/llmagent/state_agent_test.go index fef38ba37..9256e9dca 100644 --- a/agent/llmagent/state_agent_test.go +++ b/agent/llmagent/state_agent_test.go @@ -263,7 +263,7 @@ type WeatherResult struct { Timestamp time.Time `json:"timestamp"` } -func GetWeather(ctx tool.Context, args WeatherArgs) (WeatherResult, error) { +func GetWeather(ctx agent.ToolContext, args WeatherArgs) (WeatherResult, error) { // Simulate weather data temperatures := []int{-10, -5, 0, 5, 10, 15, 20, 25, 30, 35} conditions := []string{"sunny", "cloudy", "rainy", "snowy", "windy"} @@ -291,7 +291,7 @@ type CalculationResult struct { Timestamp time.Time `json:"timestamp"` } -func Calculate(ctx tool.Context, args CalculationArgs) (CalculationResult, error) { +func Calculate(ctx agent.ToolContext, args CalculationArgs) (CalculationResult, error) { operations := map[string]float64{ "add": args.X + args.Y, "subtract": args.X - args.Y, @@ -341,7 +341,7 @@ type LogActivityResult struct { err error } -func LogActivity(ctx tool.Context, params LogActivityParams) (LogActivityResult, error) { +func LogActivity(ctx agent.ToolContext, params LogActivityParams) (LogActivityResult, error) { var activityLog []LogEntry val, err := ctx.State().Get("activity_log") if err == nil { @@ -366,7 +366,7 @@ func LogActivity(ctx tool.Context, params LogActivityParams) (LogActivityResult, // --- Before Tool Callbacks --- -func beforeToolAuditCallback(ctx tool.Context, t tool.Tool, args map[string]any) (map[string]any, error) { +func beforeToolAuditCallback(ctx agent.ToolContext, t tool.Tool, args map[string]any) (map[string]any, error) { fmt.Printf("🔍 AUDIT: About to call tool '%s' with args: %v\n", t.Name(), args) var auditLog []map[string]any @@ -387,7 +387,7 @@ func beforeToolAuditCallback(ctx tool.Context, t tool.Tool, args map[string]any) return nil, nil // Continue execution } -func beforeToolSecurityCallback(ctx tool.Context, t tool.Tool, args map[string]any) (map[string]any, error) { +func beforeToolSecurityCallback(ctx agent.ToolContext, t tool.Tool, args map[string]any) (map[string]any, error) { if t.Name() == "get_weather" { location := "" if loc, ok := args["location"].(string); ok { @@ -411,7 +411,7 @@ func beforeToolSecurityCallback(ctx tool.Context, t tool.Tool, args map[string]a return nil, nil // Continue execution } -func beforeToolValidationCallback(ctx tool.Context, t tool.Tool, args map[string]any) (map[string]any, error) { +func beforeToolValidationCallback(ctx agent.ToolContext, t tool.Tool, args map[string]any) (map[string]any, error) { if t.Name() == "calculate" { operation, _ := args["operation"].(string) y, yOK := args["y"].(float64) @@ -430,7 +430,7 @@ func beforeToolValidationCallback(ctx tool.Context, t tool.Tool, args map[string // --- After Tool Callbacks --- -func afterToolEnhancementCallback(ctx tool.Context, t tool.Tool, args, result map[string]any, err error) (map[string]any, error) { +func afterToolEnhancementCallback(ctx agent.ToolContext, t tool.Tool, args, result map[string]any, err error) (map[string]any, error) { if err != nil { return result, err // Don't enhance if there was an error } @@ -444,7 +444,7 @@ func afterToolEnhancementCallback(ctx tool.Context, t tool.Tool, args, result ma return enhancedResponse, nil } -func afterToolAsyncCallback(ctx tool.Context, t tool.Tool, args, result map[string]any, err error) (map[string]any, error) { +func afterToolAsyncCallback(ctx agent.ToolContext, t tool.Tool, args, result map[string]any, err error) (map[string]any, error) { if err != nil { return result, err } @@ -653,7 +653,7 @@ func TestToolCallbacksAgent(t *testing.T) { type mockToolset struct{} -func (m *mockToolset) ProcessRequest(ctx tool.Context, req *model.LLMRequest) error { +func (m *mockToolset) ProcessRequest(ctx agent.ToolContext, req *model.LLMRequest) error { utils.AppendInstructions(req, "Extra instruction from mockToolset") return nil } diff --git a/agent/remoteagent/a2a_agent.go b/agent/remoteagent/a2a_agent.go index 6e9950052..593e840bd 100644 --- a/agent/remoteagent/a2a_agent.go +++ b/agent/remoteagent/a2a_agent.go @@ -12,22 +12,29 @@ // See the License for the specific language governing permissions and // limitations under the License. +// Package remoteagent allows using a remote ADK agents. +// +// Deprecated: Use google.golang.org/adk/agent/remoteagent/v2 instead. package remoteagent import ( "context" + "errors" "fmt" "iter" - "log" - "time" "github.com/a2aproject/a2a-go/a2a" "github.com/a2aproject/a2a-go/a2aclient" "github.com/a2aproject/a2a-go/a2aclient/agentcard" + "github.com/a2aproject/a2a-go/log" + a2av2 "github.com/a2aproject/a2a-go/v2/a2a" + a2aclientv2 "github.com/a2aproject/a2a-go/v2/a2aclient" + "github.com/a2aproject/a2a-go/v2/a2acompat/a2av0" + + "google.golang.org/genai" "google.golang.org/adk/agent" - agentinternal "google.golang.org/adk/internal/agent" - iremoteagent "google.golang.org/adk/internal/agent/remoteagent" + v2 "google.golang.org/adk/agent/remoteagent/v2" "google.golang.org/adk/server/adka2a" "google.golang.org/adk/session" ) @@ -125,236 +132,210 @@ func NewA2A(cfg A2AConfig) (agent.Agent, error) { return nil, fmt.Errorf("either AgentCard or AgentCardSource must be provided") } - remoteAgent := &a2aAgent{ - serverConfig: &iremoteagent.A2AServerConfig{ - AgentCard: cfg.AgentCard, - AgentCardSource: cfg.AgentCardSource, - CardResolveOptions: cfg.CardResolveOptions, - ClientFactory: cfg.ClientFactory, - }, - } - agent, err := agent.New(agent.Config{ + v1Cfg := v2.A2AConfig{ Name: cfg.Name, Description: cfg.Description, BeforeAgentCallbacks: cfg.BeforeAgentCallbacks, AfterAgentCallbacks: cfg.AfterAgentCallbacks, - Run: func(ic agent.InvocationContext) iter.Seq2[*session.Event, error] { - return remoteAgent.run(ic, cfg) - }, - }) - if err != nil { - return nil, err - } + ClientProvider: func(ctx context.Context, card *a2av2.AgentCard) (v2.A2AClient, error) { + if cfg.ClientFactory == nil { + factory := a2aclientv2.NewFactory( + a2aclientv2.WithCompatTransport( + a2av0.Version, a2av2.TransportProtocolJSONRPC, a2av0.NewJSONRPCTransportFactory(a2av0.JSONRPCTransportConfig{}), + ), + a2aclientv2.WithCompatTransport( + a2av0.Version, a2av2.TransportProtocolHTTPJSON, a2av0.NewRESTTransportFactory(a2av0.RESTTransportConfig{}), + ), + ) + return factory.CreateFromCard(ctx, card) + } - internalAgent, ok := agent.(agentinternal.Agent) - if !ok { - return nil, fmt.Errorf("internal error: failed to convert to internal agent") + legacyCard := a2av0.FromV1AgentCard(card) + client, err := cfg.ClientFactory.CreateFromCard(ctx, legacyCard) + if err != nil { + return nil, err + } + return &compatClient{client: client}, nil + }, } - state := agentinternal.Reveal(internalAgent) - state.AgentType = agentinternal.TypeRemoteAgent - state.Config = iremoteagent.RemoteAgentState{A2A: remoteAgent.serverConfig} - - return agent, nil -} - -type a2aAgent struct { - serverConfig *iremoteagent.A2AServerConfig -} -func (a *a2aAgent) run(ctx agent.InvocationContext, cfg A2AConfig) iter.Seq2[*session.Event, error] { - return func(yield func(*session.Event, error) bool) { - agentCard, client, err := iremoteagent.CreateA2AClient(ctx, a.serverConfig) - if err != nil { - yield(toErrorEvent(ctx, fmt.Errorf("client creation failed: %w", err)), nil) - return + if cfg.AgentCard != nil { + v1Cfg.AgentCard = a2av0.ToV1AgentCard(cfg.AgentCard) + } else if cfg.AgentCardSource != "" { + source := cfg.AgentCardSource + resolveOpts := cfg.CardResolveOptions + v1Cfg.AgentCardProvider = func(ctx context.Context) (*a2av2.AgentCard, error) { + v0Card, err := agentcard.DefaultResolver.Resolve(ctx, source, resolveOpts...) + if err != nil { + return nil, err + } + return a2av0.ToV1AgentCard(v0Card), nil } - defer destroy(client) + } - msg, err := newMessage(ctx, cfg) + if cfg.MessageSendConfig != nil { + req, err := a2av0.ToV1SendMessageRequest(&a2a.MessageSendParams{Config: cfg.MessageSendConfig}) if err != nil { - yield(toErrorEvent(ctx, fmt.Errorf("message creation failed: %w", err)), nil) - return + return nil, fmt.Errorf("MessageSendConfig conversion failed: %w", err) } + v1Cfg.MessageSendConfig = req.Config + } - req := &a2a.MessageSendParams{Message: msg, Config: cfg.MessageSendConfig} - - processor := newRunProcessor(cfg, req) - - if bcbResp, bcbErr := processor.runBeforeA2ARequestCallbacks(ctx); bcbResp != nil || bcbErr != nil { - if acbResp, acbErr := processor.runAfterA2ARequestCallbacks(ctx, bcbResp, bcbErr); acbResp != nil || acbErr != nil { - yield(acbResp, acbErr) - } else { - yield(bcbResp, bcbErr) + if cfg.Converter != nil { + v1Cfg.Converter = func(ctx agent.InvocationContext, req *a2av2.SendMessageRequest, event a2av2.Event, err error) (*session.Event, error) { + legacyReq := a2av0.FromV1SendMessageRequest(req) + var legacyEvent a2a.Event + if event != nil { + var convErr error + legacyEvent, convErr = a2av0.FromV1Event(event) + if convErr != nil { + return nil, errors.Join(fmt.Errorf("a2a event conversion failed: %w", convErr), err) + } } - return + return cfg.Converter(ctx, legacyReq, legacyEvent, err) } + } - if len(msg.Parts) == 0 { - resp := adka2a.NewRemoteAgentEvent(ctx) - if cbResp, cbErr := processor.runAfterA2ARequestCallbacks(ctx, resp, err); cbResp != nil || cbErr != nil { - yield(cbResp, cbErr) - } else { - yield(resp, nil) - } - return + if cfg.BeforeRequestCallbacks != nil { + v1Cfg.BeforeRequestCallbacks = make([]v2.BeforeA2ARequestCallback, 0, len(cfg.BeforeRequestCallbacks)) + for _, cb := range cfg.BeforeRequestCallbacks { + v1Cfg.BeforeRequestCallbacks = append(v1Cfg.BeforeRequestCallbacks, func(ctx agent.CallbackContext, req *a2av2.SendMessageRequest) (*session.Event, error) { + legacyReq := a2av0.FromV1SendMessageRequest(req) + resp, err := cb(ctx, legacyReq) + if resp != nil || err != nil { // short-circuit, no need to convert the request back + return resp, err + } + // callback pass-through request modifications + v1Req, convErr := a2av0.ToV1SendMessageRequest(legacyReq) + if convErr != nil { + return nil, convErr + } + *req = *v1Req + return nil, nil + }) } + } - var lastErr error - yieldErr := func(err error) bool { - lastErr = err - return yield(nil, err) + if cfg.AfterRequestCallbacks != nil { + v1Cfg.AfterRequestCallbacks = make([]v2.AfterA2ARequestCallback, 0, len(cfg.AfterRequestCallbacks)) + for _, cb := range cfg.AfterRequestCallbacks { + v1Cfg.AfterRequestCallbacks = append(v1Cfg.AfterRequestCallbacks, func(ctx agent.CallbackContext, req *a2av2.SendMessageRequest, resp *session.Event, err error) (*session.Event, error) { + legacyReq := a2av0.FromV1SendMessageRequest(req) + newResp, newErr := cb(ctx, legacyReq, resp, err) + if newResp != nil || newErr != nil { // short-circuit, no need to convert the request back + return newResp, newErr + } + // callback pass-through request modifications + v1Req, convErr := a2av0.ToV1SendMessageRequest(legacyReq) + if convErr != nil { + return nil, convErr + } + *req = *v1Req + return nil, nil + }) } + } - var lastEvent a2a.Event - defer func() { - err := lastErr - if err == nil && ctx.Err() != nil { - err = context.Cause(ctx) + if cfg.A2APartConverter != nil { + v1Cfg.A2APartConverter = func(ctx context.Context, a2aEvent a2av2.Event, part *a2av2.Part) (*genai.Part, error) { + legacyEvent, convErr := a2av0.FromV1Event(a2aEvent) + if convErr != nil { + return nil, convErr } - cleanupRemoteTask(ctx, cfg, agentCard, client, lastEvent, err) - }() + return cfg.A2APartConverter(ctx, legacyEvent, a2av0.FromV1Part(part)) + } + } - processEvent := func(a2aEvent a2a.Event, a2aErr error) bool { - if a2aEvent != nil { - lastEvent = a2aEvent + if cfg.GenAIPartConverter != nil { + v1Cfg.GenAIPartConverter = func(ctx context.Context, adkEvent *session.Event, part *genai.Part) (*a2av2.Part, error) { + legacyPart, err := cfg.GenAIPartConverter(ctx, adkEvent, part) + if err != nil { + return nil, err } + return a2av0.ToV1Part(legacyPart), nil + } + } - var err error - var event *session.Event - if cfg.Converter != nil { - event, err = cfg.Converter(ctx, req, a2aEvent, a2aErr) - } else { - event, err = processor.convertToSessionEvent(ctx, a2aEvent, a2aErr) - } + if cfg.RemoteTaskCleanupCallback != nil { + v1Cfg.RemoteTaskCleanupCallback = func(ctx context.Context, card *a2av2.AgentCard, client v2.A2AClient, taskInfo a2av2.TaskInfo, cause error) { + legacyCard := a2av0.FromV1AgentCard(card) + legacyTaskInfo := a2a.TaskInfo{TaskID: a2a.TaskID(taskInfo.TaskID), ContextID: taskInfo.ContextID} - if cbResp, cbErr := processor.runAfterA2ARequestCallbacks(ctx, event, err); cbResp != nil || cbErr != nil { - if cbErr != nil { - return yieldErr(cbErr) - } - event = cbResp - err = nil + if cc, ok := client.(*compatClient); ok { + cfg.RemoteTaskCleanupCallback(ctx, legacyCard, cc.client, legacyTaskInfo, cause) + return } - if err != nil { - return yieldErr(err) - } + log.Warn(ctx, "client is not an instance of compatClient, fallback to creating a new client", "type", fmt.Sprintf("%T", client)) - if event != nil { // an event might be skipped - for _, toEmit := range processor.aggregatePartial(ctx, a2aEvent, event) { - if !yield(toEmit, nil) { - return false - } - } + factory := cfg.ClientFactory + if factory == nil { + factory = a2aclient.NewFactory() } - return true - } - - if ctx.RunConfig().StreamingMode == agent.StreamingModeNone { - a2aEvent, a2aErr := client.SendMessage(ctx, req) - processEvent(a2aEvent, a2aErr) - return - } - - for a2aEvent, a2aErr := range client.SendStreamingMessage(ctx, req) { - if !processEvent(a2aEvent, a2aErr) { + legacyClient, err := factory.CreateFromCard(ctx, legacyCard) + if err != nil { + log.Warn(ctx, "RemoteTaskCleanupCallback: failed to create legacy client", "error", err) return } + defer func() { + if err := legacyClient.Destroy(); err != nil { + log.Warn(ctx, "RemoteTaskCleanupCallback: failed to destroy a legacy client", "error", err) + } + }() + cfg.RemoteTaskCleanupCallback(ctx, legacyCard, legacyClient, legacyTaskInfo, cause) } } -} -func cleanupRemoteTask(ctx context.Context, cfg A2AConfig, card *a2a.AgentCard, client *a2aclient.Client, lastEvent a2a.Event, cause error) { - if lastEvent == nil { - return - } - taskID := lastEvent.TaskInfo().TaskID - if taskID == "" { - return - } - if _, ok := lastEvent.(*a2a.Message); ok { - return - } - var state a2a.TaskState - if tu, ok := lastEvent.(*a2a.TaskStatusUpdateEvent); ok { - state = tu.Status.State - } - if t, ok := lastEvent.(*a2a.Task); ok { - state = t.Status.State - } - if state.Terminal() { - return - } - - ctx = context.WithoutCancel(ctx) + return v2.NewA2A(v1Cfg) +} - if cfg.RemoteTaskCleanupCallback != nil { - cfg.RemoteTaskCleanupCallback(ctx, card, client, lastEvent.TaskInfo(), cause) - return - } +type compatClient struct { + client *a2aclient.Client +} - if state == a2a.TaskStateInputRequired && cause == nil { - return +func (s *compatClient) SendMessage(ctx context.Context, req *a2av2.SendMessageRequest) (a2av2.SendMessageResult, error) { + legacyResp, err := s.client.SendMessage(ctx, a2av0.FromV1SendMessageRequest(req)) + if err != nil { + return nil, err } - cancelCtx, cancelTimeout := context.WithTimeout(ctx, 5*time.Second) - defer cancelTimeout() - _, err := client.CancelTask(cancelCtx, &a2a.TaskIDParams{ID: taskID}) + v1Event, err := a2av0.ToV1Event(legacyResp) if err != nil { - log.Printf("failed to cancel task %s: %v", taskID, err) + return nil, err } -} - -func newMessage(ctx agent.InvocationContext, cfg A2AConfig) (*a2a.Message, error) { - events := ctx.Session().Events() - if userFnCall := getUserFunctionCallAt(events, events.Len()-1); userFnCall != nil { - event := userFnCall.response - parts, err := convertParts(ctx, cfg, event) - if err != nil { - return nil, err - } - msg := a2a.NewMessage(a2a.MessageRoleUser, parts...) - msg.TaskID = userFnCall.taskID - msg.ContextID = userFnCall.contextID - return msg, nil + res, ok := v1Event.(a2av2.SendMessageResult) + if !ok { + return nil, fmt.Errorf("converted event does not implement SendMessageResult: %T", v1Event) } - - parts, contextID := toMissingRemoteSessionParts(ctx, events, cfg) - msg := a2a.NewMessage(a2a.MessageRoleUser, parts...) - msg.ContextID = contextID - return msg, nil + return res, nil } -func toErrorEvent(ctx agent.InvocationContext, err error) *session.Event { - event := adka2a.NewRemoteAgentEvent(ctx) - event.ErrorMessage = err.Error() - event.CustomMetadata = map[string]any{adka2a.ToADKMetaKey("error"): err.Error()} - event.TurnComplete = true - return event -} - -func convertParts(ctx agent.InvocationContext, cfg A2AConfig, event *session.Event) ([]a2a.Part, error) { - parts := make([]a2a.Part, 0, len(event.Content.Parts)) - if cfg.GenAIPartConverter != nil { - for _, part := range event.Content.Parts { - cp, err := cfg.GenAIPartConverter(ctx, event, part) +func (s *compatClient) SendStreamingMessage(ctx context.Context, req *a2av2.SendMessageRequest) iter.Seq2[a2av2.Event, error] { + return func(yield func(a2av2.Event, error) bool) { + for legacyEvent, err := range s.client.SendStreamingMessage(ctx, a2av0.FromV1SendMessageRequest(req)) { if err != nil { - return nil, err + yield(nil, err) + return } - if cp != nil { - parts = append(parts, cp) + v1Event, convErr := a2av0.ToV1Event(legacyEvent) + if convErr != nil { + yield(nil, convErr) + return + } + if !yield(v1Event, nil) { + return } - } - } else { - var err error - parts, err = adka2a.ToA2AParts(event.Content.Parts, event.LongRunningToolIDs) - if err != nil { - return nil, fmt.Errorf("event part conversion failed: %w", err) } } - return parts, nil } -func destroy(client *a2aclient.Client) { - if err := client.Destroy(); err != nil { - log.Printf("failed to destroy client: %v", err) +func (s *compatClient) CancelTask(ctx context.Context, req *a2av2.CancelTaskRequest) (*a2av2.Task, error) { + legacyResp, err := s.client.CancelTask(ctx, a2av0.FromV1CancelTaskRequest(req)) + if err != nil { + return nil, err } + return a2av0.ToV1Task(legacyResp) +} + +func (s *compatClient) Destroy() error { + return s.client.Destroy() } diff --git a/agent/remoteagent/a2a_agent_compat_test.go b/agent/remoteagent/a2a_agent_compat_test.go new file mode 100644 index 000000000..e753dcea9 --- /dev/null +++ b/agent/remoteagent/a2a_agent_compat_test.go @@ -0,0 +1,810 @@ +// Copyright 2025 Google LLC +// +// Licensed under the Apache License, Version 2.0 (the "License"); +// you may not use this file except in compliance with the License. +// You may obtain a copy of the License at +// +// http://www.apache.org/licenses/LICENSE-2.0 +// +// Unless required by applicable law or agreed to in writing, software +// distributed under the License is distributed on an "AS IS" BASIS, +// WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. +// See the License for the specific language governing permissions and +// limitations under the License. + +package remoteagent + +import ( + "context" + "iter" + "net/http/httptest" + "testing" + "time" + + legacyA2A "github.com/a2aproject/a2a-go/a2a" + legacyAClient "github.com/a2aproject/a2a-go/a2aclient" + legacyASrv "github.com/a2aproject/a2a-go/a2asrv" + legacyEQ "github.com/a2aproject/a2a-go/a2asrv/eventqueue" + v2a2a "github.com/a2aproject/a2a-go/v2/a2a" + "github.com/a2aproject/a2a-go/v2/a2acompat/a2av0" + v2asrv "github.com/a2aproject/a2a-go/v2/a2asrv" + "google.golang.org/genai" + + "google.golang.org/adk/agent" + "google.golang.org/adk/agent/llmagent" + icontext "google.golang.org/adk/internal/context" + "google.golang.org/adk/internal/utils" + "google.golang.org/adk/model" + "google.golang.org/adk/runner" + "google.golang.org/adk/server/adka2a" + "google.golang.org/adk/session" +) + +func TestCompat_OldExecutor_Direct(t *testing.T) { + agentName := "test-agent" + agentObj := utils.Must(agent.New(agent.Config{ + Name: agentName, + Run: func(ic agent.InvocationContext) iter.Seq2[*session.Event, error] { + return func(yield func(*session.Event, error) bool) { + yield(&session.Event{ + LLMResponse: model.LLMResponse{Content: genai.NewContentFromText("hello", genai.RoleModel)}, + Author: agentName, + }, nil) + } + }, + })) + + executor := adka2a.NewExecutor(adka2a.ExecutorConfig{ + OutputMode: adka2a.OutputArtifactPerEvent, + RunnerConfig: runner.Config{ + AppName: "TestApp", + Agent: agentObj, + SessionService: session.InMemoryService(), + }, + RunConfig: agent.RunConfig{ + StreamingMode: agent.StreamingModeSSE, + }, + GenAIPartConverter: func(ctx context.Context, adkEvent *session.Event, part *genai.Part) (legacyA2A.Part, error) { + return a2av0.FromV1Part(v2a2a.NewTextPart(part.Text)), nil + }, + AfterEventCallback: func(ctx adka2a.ExecutorContext, event *session.Event, processed *legacyA2A.TaskArtifactUpdateEvent) error { + if processed.Artifact != nil && len(processed.Artifact.Parts) > 0 { + processed.Artifact.Parts[0] = a2av0.FromV1Part(v2a2a.NewTextPart("modified-by-executor")) + } + return nil + }, + }) + + reqCtx := &legacyASrv.RequestContext{ + ContextID: "test-context", + TaskID: legacyA2A.NewTaskID(), + Message: legacyA2A.NewMessage(legacyA2A.MessageRoleUser, a2av0.FromV1Part(v2a2a.NewTextPart("hi"))), + } + queue := &mockQueue{} + if err := executor.Execute(t.Context(), reqCtx, queue); err != nil { + t.Fatalf("executor.Execute() error = %v", err) + } + + found := false + for _, ev := range queue.events { + if ae, ok := ev.(*legacyA2A.TaskArtifactUpdateEvent); ok { + for _, p := range ae.Artifact.Parts { + gp, _ := adka2a.ToGenAIPart(p) + if gp != nil && gp.Text == "modified-by-executor" { + found = true + } + } + } + } + if !found { + t.Error("Did not find modified part in executor output events") + } +} + +func TestCompat_RemoteAgent(t *testing.T) { + tests := []struct { + name string + executor *mockV2Executor + updateConfig func(config *A2AConfig) + wantEventWithText string + }{ + { + name: "after request callback modifies response", + executor: &mockV2Executor{ + events: []v2a2a.Event{ + v2a2a.NewMessage(v2a2a.MessageRoleAgent, v2a2a.NewTextPart("hello")), + }, + }, + updateConfig: func(config *A2AConfig) { + config.AfterRequestCallbacks = []AfterA2ARequestCallback{ + func(ctx agent.CallbackContext, req *legacyA2A.MessageSendParams, resp *session.Event, err error) (*session.Event, error) { + if resp != nil && resp.Content != nil && len(resp.Content.Parts) > 0 { + resp.Content.Parts[0].Text = "modified-by-agent-callback" + } + return nil, nil + }, + } + }, + wantEventWithText: "modified-by-agent-callback", + }, + { + name: "before request callback modifies request", + executor: &mockV2Executor{ + executeFn: func(ctx context.Context, execCtx *v2asrv.ExecutorContext) iter.Seq2[v2a2a.Event, error] { + return func(yield func(v2a2a.Event, error) bool) { + text := "" + if execCtx.Message != nil && len(execCtx.Message.Parts) > 0 { + text = execCtx.Message.Parts[0].Text() + } + yield(v2a2a.NewMessage(v2a2a.MessageRoleAgent, v2a2a.NewTextPart("echo:"+text)), nil) + } + }, + }, + updateConfig: func(config *A2AConfig) { + config.BeforeRequestCallbacks = []BeforeA2ARequestCallback{ + func(ctx agent.CallbackContext, req *legacyA2A.MessageSendParams) (*session.Event, error) { + req.Message = legacyA2A.NewMessage(legacyA2A.MessageRoleUser, legacyA2A.TextPart{Text: "42"}) + return nil, nil + }, + } + }, + wantEventWithText: "echo:42", + }, + { + name: "before request callback short-circuits", + executor: &mockV2Executor{ + executeFn: func(ctx context.Context, execCtx *v2asrv.ExecutorContext) iter.Seq2[v2a2a.Event, error] { + return func(yield func(v2a2a.Event, error) bool) { + t.Fatal("server should not be called when before callback short-circuits") + } + }, + }, + updateConfig: func(config *A2AConfig) { + config.BeforeRequestCallbacks = []BeforeA2ARequestCallback{ + func(ctx agent.CallbackContext, req *legacyA2A.MessageSendParams) (*session.Event, error) { + return &session.Event{ + LLMResponse: model.LLMResponse{ + Content: genai.NewContentFromText("cached-response", genai.RoleModel), + }, + }, nil + }, + } + }, + wantEventWithText: "cached-response", + }, + { + name: "custom converter", + executor: &mockV2Executor{ + events: []v2a2a.Event{ + v2a2a.NewMessage(v2a2a.MessageRoleAgent, v2a2a.NewTextPart("original")), + }, + }, + updateConfig: func(config *A2AConfig) { + config.Converter = func(ctx agent.InvocationContext, req *legacyA2A.MessageSendParams, event legacyA2A.Event, err error) (*session.Event, error) { + ev := session.NewEventWithContext(ctx, "custom") + ev.Author = ctx.Agent().Name() + ev.LLMResponse = model.LLMResponse{ + Content: genai.NewContentFromText("converted", genai.RoleModel), + TurnComplete: true, + } + return ev, nil + } + }, + wantEventWithText: "converted", + }, + { + name: "GenAI part converter", + executor: &mockV2Executor{ + executeFn: func(ctx context.Context, execCtx *v2asrv.ExecutorContext) iter.Seq2[v2a2a.Event, error] { + return func(yield func(v2a2a.Event, error) bool) { + if execCtx.Message != nil { + for _, p := range execCtx.Message.Parts { + if p.Text() == "custom:hello" { + yield(v2a2a.NewMessage(v2a2a.MessageRoleAgent, v2a2a.NewTextPart("converter-verified")), nil) + return + } + } + } + yield(v2a2a.NewMessage(v2a2a.MessageRoleAgent, v2a2a.NewTextPart("converter-not-applied")), nil) + } + }, + }, + updateConfig: func(config *A2AConfig) { + config.GenAIPartConverter = func(ctx context.Context, adkEvent *session.Event, part *genai.Part) (legacyA2A.Part, error) { + if part.Text != "" { + return a2av0.FromV1Part(v2a2a.NewTextPart("custom:" + part.Text)), nil + } + return adka2a.ToA2APart(part, nil) + } + }, + wantEventWithText: "converter-verified", + }, + { + name: "A2A part converter", + executor: &mockV2Executor{ + events: []v2a2a.Event{ + v2a2a.NewMessage(v2a2a.MessageRoleAgent, v2a2a.NewTextPart("raw-response")), + }, + }, + updateConfig: func(config *A2AConfig) { + config.A2APartConverter = func(ctx context.Context, a2aEvent legacyA2A.Event, part legacyA2A.Part) (*genai.Part, error) { + tp, ok := part.(legacyA2A.TextPart) + if ok { + return genai.NewPartFromText("custom:" + tp.Text), nil + } + return adka2a.ToGenAIPart(part) + } + }, + wantEventWithText: "custom:raw-response", + }, + { + name: "multiple after request callbacks execute in order", + executor: &mockV2Executor{ + events: []v2a2a.Event{ + v2a2a.NewMessage(v2a2a.MessageRoleAgent, v2a2a.NewTextPart("hello")), + }, + }, + updateConfig: func(config *A2AConfig) { + config.AfterRequestCallbacks = []AfterA2ARequestCallback{ + func(ctx agent.CallbackContext, req *legacyA2A.MessageSendParams, resp *session.Event, err error) (*session.Event, error) { + if resp != nil && resp.Content != nil && len(resp.Content.Parts) > 0 { + resp.Content.Parts[0].Text += "-first" + } + return nil, nil + }, + func(ctx agent.CallbackContext, req *legacyA2A.MessageSendParams, resp *session.Event, err error) (*session.Event, error) { + if resp != nil && resp.Content != nil && len(resp.Content.Parts) > 0 { + resp.Content.Parts[0].Text += "-second" + } + return nil, nil + }, + } + }, + wantEventWithText: "hello-first-second", + }, + } + + for _, tc := range tests { + t.Run(tc.name, func(t *testing.T) { + server := startA2AServer(t, tc.executor) + card := newLegacyCard(server.URL) + config := A2AConfig{Name: "remote-agent", AgentCard: card} + tc.updateConfig(&config) + agnt, err := NewA2A(config) + if err != nil { + t.Fatalf("NewA2A() error = %v", err) + } + ic := newInvocationContext(t, []*session.Event{newUserHello()}) + events, err := runAndCollect(ic, agnt) + if err != nil { + t.Fatalf("agent.Run() error = %v", err) + } + foundText := false + var texts []string + for _, ev := range events { + if foundText { + break + } + if ev.Content == nil { + continue + } + for _, p := range ev.Content.Parts { + if p.Text == tc.wantEventWithText { + foundText = true + break + } + if p.Text != "" { + texts = append(texts, p.Text) + } + } + } + if !foundText { + t.Errorf("expected text %q in events, got texts: %v", tc.wantEventWithText, texts) + } + }) + } +} + +func TestCompat_RemoteTaskCleanupCallback(t *testing.T) { + mockExec := &mockV2Executor{ + executeFn: func(ctx context.Context, execCtx *v2asrv.ExecutorContext) iter.Seq2[v2a2a.Event, error] { + return func(yield func(v2a2a.Event, error) bool) { + if !yield(v2a2a.NewSubmittedTask(execCtx, execCtx.Message), nil) { + return + } + for ctx.Err() == nil { + data := v2a2a.NewDataPart(map[string]any{"tick": true}) + if !yield(v2a2a.NewArtifactEvent(execCtx, data), nil) { + return + } + time.Sleep(1 * time.Millisecond) + } + yield(v2a2a.NewStatusUpdateEvent(execCtx, v2a2a.TaskStateCompleted, nil), nil) + } + }, + } + + server := startA2AServer(t, mockExec) + card := newLegacyCard(server.URL) + + cleanupCalled := false + var cleanupTaskInfo legacyA2A.TaskInfo + oldAgent := utils.Must(NewA2A(A2AConfig{ + Name: "remote-agent", + AgentCard: card, + RemoteTaskCleanupCallback: func(ctx context.Context, card *legacyA2A.AgentCard, client *legacyAClient.Client, taskInfo legacyA2A.TaskInfo, cause error) { + cleanupCalled = true + cleanupTaskInfo = taskInfo + }, + })) + + // Use a cancelable context so we can trigger cleanup by canceling mid-stream. + ctx, cancel := context.WithCancel(t.Context()) + defer cancel() + svc := session.InMemoryService() + resp, err := svc.Create(ctx, &session.CreateRequest{AppName: t.Name(), UserID: "test"}) + if err != nil { + t.Fatalf("session.Create() error = %v", err) + } + hello := newUserHello() + if err := svc.AppendEvent(ctx, resp.Session, hello); err != nil { + t.Fatalf("AppendEvent() error = %v", err) + } + ic := icontext.NewInvocationContext(ctx, icontext.InvocationContextParams{ + Session: resp.Session, + RunConfig: &agent.RunConfig{StreamingMode: agent.StreamingModeSSE}, + }) + + // Break out of the run after receiving a couple events to trigger cleanup. + count := 0 + for _, err := range oldAgent.Run(ic) { + if err != nil { + break + } + count++ + if count >= 2 { + cancel() + } + } + + if !cleanupCalled { + t.Error("RemoteTaskCleanupCallback was not called") + } + if cleanupTaskInfo.TaskID == "" { + t.Error("RemoteTaskCleanupCallback received empty TaskID") + } +} + +func TestCompat_ContextPropagation(t *testing.T) { + appName := "yesapp" + sessionService := session.InMemoryService() + agnt := utils.Must(agent.New(agent.Config{ + Name: appName, + Run: func(ic agent.InvocationContext) iter.Seq2[*session.Event, error] { + return func(yield func(*session.Event, error) bool) { + yield(&session.Event{LLMResponse: model.LLMResponse{Content: genai.NewContentFromText("Yes", genai.RoleModel)}}, nil) + } + }, + })) + executor := adka2a.NewExecutor(adka2a.ExecutorConfig{ + RunnerConfig: runner.Config{ + AppName: appName, + Agent: agnt, + SessionService: sessionService, + }, + }) + + handler := legacyASrv.NewHandler( + executor, + legacyASrv.WithCallInterceptor(&testSrvAuthInterceptor{ + BeforeFunc: func(ctx context.Context, callCtx *legacyASrv.CallContext, req *legacyASrv.Request) (context.Context, error) { + if headers, _ := callCtx.RequestMeta().Get("authorization"); len(headers) > 0 { + callCtx.User = &legacyASrv.AuthenticatedUser{UserName: headers[0]} + } + return ctx, nil + }, + }), + ) + server := httptest.NewServer(legacyASrv.NewJSONRPCHandler(handler)) + t.Cleanup(server.Close) + + userID := "user123" + remote, err := NewA2A(A2AConfig{ + Name: "yes-client", + AgentCard: newLegacyCard(server.URL), + ClientFactory: legacyAClient.NewFactory(legacyAClient.WithInterceptors( + &testClientAuthInterceptor{ + BeforeFunc: func(ctx context.Context, req *legacyAClient.Request) (context.Context, error) { + req.Meta["Authorization"] = []string{userID} + return ctx, nil + }, + }, + )), + }) + if err != nil { + t.Fatalf("NewA2A() error = %v", err) + } + _, err = runAndCollect(newInvocationContext(t, []*session.Event{newUserHello()}), remote) + if err != nil { + t.Fatalf("agent.Run() error = %v", err) + } + + sessions, err := sessionService.List(t.Context(), &session.ListRequest{ + AppName: appName, + UserID: userID, + }) + if err != nil { + t.Fatalf("sessionService.List() error = %v", err) + } + if len(sessions.Sessions) != 1 { + t.Fatalf("len(sessions.Sessions) = %d, want 1", len(sessions.Sessions)) + } +} + +type testSrvAuthInterceptor struct { + legacyASrv.PassthroughCallInterceptor + BeforeFunc func(ctx context.Context, callCtx *legacyASrv.CallContext, req *legacyASrv.Request) (context.Context, error) +} + +func (i *testSrvAuthInterceptor) Before(ctx context.Context, callCtx *legacyASrv.CallContext, req *legacyASrv.Request) (context.Context, error) { + if i.BeforeFunc != nil { + return i.BeforeFunc(ctx, callCtx, req) + } + return ctx, nil +} + +type testClientAuthInterceptor struct { + legacyAClient.PassthroughInterceptor + BeforeFunc func(ctx context.Context, req *legacyAClient.Request) (context.Context, error) +} + +func (i *testClientAuthInterceptor) Before(ctx context.Context, req *legacyAClient.Request) (context.Context, error) { + if i.BeforeFunc != nil { + return i.BeforeFunc(ctx, req) + } + return ctx, nil +} + +type mockQueue struct { + events []legacyA2A.Event +} + +func (q *mockQueue) Write(ctx context.Context, event legacyA2A.Event) error { + q.events = append(q.events, event) + return nil +} + +func (q *mockQueue) WriteVersioned(ctx context.Context, event legacyA2A.Event, version legacyA2A.TaskVersion) error { + return q.Write(ctx, event) +} + +func (q *mockQueue) Read(ctx context.Context) (legacyA2A.Event, legacyA2A.TaskVersion, error) { + var v legacyA2A.TaskVersion + return nil, v, nil +} + +func (q *mockQueue) Close() error { return nil } + +type mockV2Executor struct { + events []v2a2a.Event + executeFn func(ctx context.Context, execCtx *v2asrv.ExecutorContext) iter.Seq2[v2a2a.Event, error] +} + +func (e *mockV2Executor) Execute(ctx context.Context, execCtx *v2asrv.ExecutorContext) iter.Seq2[v2a2a.Event, error] { + if e.executeFn != nil { + return e.executeFn(ctx, execCtx) + } + return func(yield func(v2a2a.Event, error) bool) { + for _, ev := range e.events { + if !yield(ev, nil) { + return + } + } + } +} + +func (e *mockV2Executor) Cancel(ctx context.Context, execCtx *v2asrv.ExecutorContext) iter.Seq2[v2a2a.Event, error] { + return func(yield func(v2a2a.Event, error) bool) { + yield(v2a2a.NewStatusUpdateEvent(execCtx, v2a2a.TaskStateCanceled, nil), nil) + } +} + +func startA2AServer(t *testing.T, executor *mockV2Executor) *httptest.Server { + t.Helper() + handler := v2asrv.NewHandler(executor) + server := httptest.NewServer(v2asrv.NewJSONRPCHandler(handler)) + t.Cleanup(server.Close) + return server +} + +func newLegacyCard(serverURL string) *legacyA2A.AgentCard { + return a2av0.FromV1AgentCard(&v2a2a.AgentCard{ + SupportedInterfaces: []*v2a2a.AgentInterface{ + v2a2a.NewAgentInterface(serverURL, v2a2a.TransportProtocolJSONRPC), + }, + Capabilities: v2a2a.AgentCapabilities{Streaming: true}, + }) +} + +func newV0Card(serverURL string) *legacyA2A.AgentCard { + return a2av0.FromV1AgentCard(&v2a2a.AgentCard{ + SupportedInterfaces: []*v2a2a.AgentInterface{ + { + URL: serverURL, + ProtocolBinding: v2a2a.TransportProtocolJSONRPC, + ProtocolVersion: a2av0.Version, + }, + }, + Capabilities: v2a2a.AgentCapabilities{Streaming: true}, + }) +} + +func newInvocationContext(t *testing.T, events []*session.Event) agent.InvocationContext { + t.Helper() + ctx := t.Context() + service := session.InMemoryService() + resp, err := service.Create(ctx, &session.CreateRequest{AppName: t.Name(), UserID: "test"}) + if err != nil { + t.Fatalf("sessionService.Create() error = %v", err) + } + for _, event := range events { + if err := service.AppendEvent(ctx, resp.Session, event); err != nil { + t.Fatalf("sessionService.AppendEvent() error = %v", err) + } + } + + ic := icontext.NewInvocationContext(ctx, icontext.InvocationContextParams{ + Session: resp.Session, + RunConfig: &agent.RunConfig{ + StreamingMode: agent.StreamingModeSSE, + }, + }) + return ic +} + +func newUserHello() *session.Event { + event := session.NewEventWithContext(context.Background(), "invocation") + event.Author = "user" + event.LLMResponse = model.LLMResponse{ + Content: genai.NewContentFromText("hello", genai.RoleUser), + } + return event +} + +func runAndCollect(ic agent.InvocationContext, agnt agent.Agent) ([]*session.Event, error) { + var collected []*session.Event + for ev, err := range agnt.Run(ic) { + if err != nil { + return collected, err + } + collected = append(collected, ev) + } + return collected, nil +} + +type mockLegacyExecutor struct { + executeFn func(ctx context.Context, reqCtx *legacyASrv.RequestContext, queue legacyEQ.Queue) error + cancelFn func(ctx context.Context, reqCtx *legacyASrv.RequestContext, queue legacyEQ.Queue) error + cleanupFn func(ctx context.Context, reqCtx *legacyASrv.RequestContext, result legacyA2A.SendMessageResult, err error) +} + +var ( + _ legacyASrv.AgentExecutor = (*mockLegacyExecutor)(nil) + _ legacyASrv.AgentExecutionCleaner = (*mockLegacyExecutor)(nil) + + transferToolName = "transfer_to_agent" + modelTextRootTransfer = "transferring... please hold... beepboop..." +) + +type testA2AServer struct { + *httptest.Server + handler legacyASrv.RequestHandler +} + +func (e *mockLegacyExecutor) Execute(ctx context.Context, reqCtx *legacyASrv.RequestContext, queue legacyEQ.Queue) error { + if e.executeFn != nil { + return e.executeFn(ctx, reqCtx, queue) + } + return nil +} + +func (e *mockLegacyExecutor) Cancel(ctx context.Context, reqCtx *legacyASrv.RequestContext, queue legacyEQ.Queue) error { + if e.cancelFn != nil { + return e.cancelFn(ctx, reqCtx, queue) + } + return queue.Write(ctx, legacyA2A.NewStatusUpdateEvent(reqCtx, legacyA2A.TaskStateCanceled, nil)) +} + +func (e *mockLegacyExecutor) Cleanup(ctx context.Context, reqCtx *legacyASrv.RequestContext, result legacyA2A.SendMessageResult, err error) { + if e.cleanupFn != nil { + e.cleanupFn(ctx, reqCtx, result, err) + } +} + +func startLegacyA2AServer(t *testing.T, executor legacyASrv.AgentExecutor) *testA2AServer { + t.Helper() + handler := legacyASrv.NewHandler(executor) + server := httptest.NewServer(legacyASrv.NewJSONRPCHandler(handler)) + return &testA2AServer{Server: server, handler: handler} +} + +func newLegacyA2AClient(t *testing.T, server *testA2AServer) *legacyAClient.Client { + t.Helper() + card := newV0Card(server.Server.URL) + client, err := legacyAClient.NewFromCard(t.Context(), card) + if err != nil { + t.Fatalf("legacyAClient.NewFromCard() error = %v", err) + } + return client +} + +func newA2ARemoteAgent(t *testing.T, name string, server *testA2AServer) agent.Agent { + t.Helper() + card := newV0Card(server.Server.URL) + + return utils.Must(NewA2A(A2AConfig{AgentCard: card, Name: name})) +} + +type llmStub struct { + name string + generateContent func(ctx context.Context, req *model.LLMRequest, stream bool) iter.Seq2[*model.LLMResponse, error] +} + +func (d *llmStub) Name() string { + return d.name +} + +func (d *llmStub) GenerateContent(ctx context.Context, req *model.LLMRequest, stream bool) iter.Seq2[*model.LLMResponse, error] { + return d.generateContent(ctx, req, stream) +} + +func newRootAgent(name string, subAgent agent.Agent) agent.Agent { + return utils.Must(llmagent.New(llmagent.Config{ + Name: name, + SubAgents: []agent.Agent{subAgent}, + Model: &llmStub{ + name: name + "-model", + generateContent: func(ctx context.Context, req *model.LLMRequest, stream bool) iter.Seq2[*model.LLMResponse, error] { + return func(yield func(*model.LLMResponse, error) bool) { + yield(&model.LLMResponse{ + Content: genai.NewContentFromParts([]*genai.Part{ + genai.NewPartFromText(modelTextRootTransfer), + genai.NewPartFromFunctionCall(transferToolName, map[string]any{"agent_name": subAgent.Name()}), + }, genai.RoleModel), + }, nil) + } + }, + }, + })) +} + +func TestCompat_A2ACleanupPropagation(t *testing.T) { + // Remote A2A server publishes a submitted task and start generating artifact updates + // until it detects a context cancelation + remoteTaskIDChan, remoteCleanupCalledChan := make(chan legacyA2A.TaskID, 1), make(chan struct{}, 2) + serverB := startLegacyA2AServer(t, &mockLegacyExecutor{ + cancelFn: func(ctx context.Context, reqCtx *legacyASrv.RequestContext, queue legacyEQ.Queue) error { + event := legacyA2A.NewStatusUpdateEvent(reqCtx, legacyA2A.TaskStateCanceled, nil) + event.Final = true + return queue.Write(ctx, event) + }, + executeFn: func(ctx context.Context, reqCtx *legacyASrv.RequestContext, queue legacyEQ.Queue) error { + remoteTaskIDChan <- reqCtx.TaskID + if err := queue.Write(ctx, legacyA2A.NewSubmittedTask(reqCtx, reqCtx.Message)); err != nil { + return err + } + for ctx.Err() == nil { + if err := queue.Write(ctx, legacyA2A.NewArtifactEvent(reqCtx, legacyA2A.TextPart{Text: "foo"})); err != nil { + return err + } + time.Sleep(1 * time.Millisecond) + } + return queue.Write(ctx, legacyA2A.NewStatusUpdateEvent(reqCtx, legacyA2A.TaskStateCompleted, nil)) + }, + cleanupFn: func(ctx context.Context, reqCtx *legacyASrv.RequestContext, result legacyA2A.SendMessageResult, cause error) { + remoteCleanupCalledChan <- struct{}{} + }, + }) + defer serverB.Close() + + // Root server connects to server B through remote subagent + remoteAgentB := newA2ARemoteAgent(t, "remote-agent-b", serverB) + rootA := newRootAgent("agent-b", remoteAgentB) + + // A2A Execution Cleanup callback + executorCleanupCalledChan := make(chan struct{}, 2) + executorA := adka2a.NewExecutor(adka2a.ExecutorConfig{ + OutputMode: adka2a.OutputArtifactPerEvent, + RunnerConfig: runner.Config{ + AppName: rootA.Name(), + SessionService: session.InMemoryService(), + Agent: rootA, + }, + A2AExecutionCleanupCallback: func(ctx context.Context, reqCtx *legacyASrv.RequestContext, subAgentCards []*legacyA2A.AgentCard, result legacyA2A.SendMessageResult, cause error) { + executorCleanupCalledChan <- struct{}{} + }, + RunConfig: agent.RunConfig{StreamingMode: agent.StreamingModeSSE}, + }) + + serverA := startLegacyA2AServer(t, executorA) + defer serverA.Close() + + client := newLegacyA2AClient(t, serverA) + + // Send a streaming message in a detached goroutine, passing status update through chan + statusUpdateEventChan := make(chan legacyA2A.Event, 10) + go func() { + defer close(statusUpdateEventChan) + msg := legacyA2A.NewMessage(legacyA2A.MessageRoleUser, legacyA2A.TextPart{Text: "work"}) + for event, err := range client.SendStreamingMessage(t.Context(), &legacyA2A.MessageSendParams{Message: msg}) { + if err != nil { + t.Errorf("client.SendStreamingMessage() error = %v", err) + return + } + if _, ok := event.(*legacyA2A.TaskArtifactUpdateEvent); ok { + continue + } + statusUpdateEventChan <- event + } + }() + + // Issue a task cancellation request + taskID := (<-statusUpdateEventChan).TaskInfo().TaskID + cancelResultChan := make(chan *legacyA2A.Task, 1) + go func() { + defer close(cancelResultChan) + task, err := client.CancelTask(t.Context(), &legacyA2A.TaskIDParams{ID: taskID}) + if err != nil { + t.Errorf("client.CancelTask() error = %v", err) + return + } + cancelResultChan <- task + }() + + // Check the streaming message sender got a cancelled state task in their response + var lastStreamingUpdate legacyA2A.Event + for event := range statusUpdateEventChan { + lastStreamingUpdate = event + } + if tu, ok := lastStreamingUpdate.(*legacyA2A.TaskStatusUpdateEvent); ok { + if tu.Status.State != legacyA2A.TaskStateCanceled { + t.Errorf("lastStreamingUpdate.Status.State = %q, want %q", tu.Status.State, legacyA2A.TaskStateCanceled) + } + } else { + t.Fatalf("type(lastStreamingUpdate) = %T, want *a2a.TaskStatusUpdateEvent", lastStreamingUpdate) + } + + // Check subagent task got cancelled when the parent task was cancelled. + // Reads from channel twice because cleanup gets called both for cancelation and execution. + timeout := time.After(5 * time.Second) + for range 2 { + select { + case <-remoteCleanupCalledChan: + case <-timeout: + t.Fatalf("remote cleanup was not called") + } + } + var remoteTaskID legacyA2A.TaskID + select { + case remoteTaskID = <-remoteTaskIDChan: + case <-time.After(1 * time.Second): + t.Fatal("server B was never reached; remoteTaskIDChan is empty") + } + + for range 2 { + select { + case <-executorCleanupCalledChan: + case <-timeout: + t.Fatalf("executor cleanup was not called") + } + } + + remoteClient := newLegacyA2AClient(t, serverB) + remoteTask, err := remoteClient.GetTask(t.Context(), &legacyA2A.TaskQueryParams{ID: remoteTaskID}) + if err != nil { + t.Fatalf("remoteClient.GetTask() error = %v", err) + } + if remoteTask.Status.State != legacyA2A.TaskStateCanceled { + t.Errorf("remoteTask.Status.State = %q, want %q", remoteTask.Status.State, legacyA2A.TaskStateCanceled) + } +} diff --git a/agent/remoteagent/v2/a2a_agent.go b/agent/remoteagent/v2/a2a_agent.go new file mode 100644 index 000000000..73c2fa5b3 --- /dev/null +++ b/agent/remoteagent/v2/a2a_agent.go @@ -0,0 +1,378 @@ +// Copyright 2025 Google LLC +// +// Licensed under the Apache License, Version 2.0 (the "License"); +// you may not use this file except in compliance with the License. +// You may obtain a copy of the License at +// +// http://www.apache.org/licenses/LICENSE-2.0 +// +// Unless required by applicable law or agreed to in writing, software +// distributed under the License is distributed on an "AS IS" BASIS, +// WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. +// See the License for the specific language governing permissions and +// limitations under the License. + +package remoteagent + +import ( + "context" + "encoding/json" + "fmt" + "iter" + "os" + "strings" + "time" + + "github.com/a2aproject/a2a-go/v2/a2a" + "github.com/a2aproject/a2a-go/v2/a2aclient" + "github.com/a2aproject/a2a-go/v2/a2aclient/agentcard" + "github.com/a2aproject/a2a-go/v2/log" + + "google.golang.org/adk/agent" + agentinternal "google.golang.org/adk/internal/agent" + iremoteagent "google.golang.org/adk/internal/agent/remoteagent" + "google.golang.org/adk/server/adka2a/v2" + "google.golang.org/adk/session" +) + +// BeforeA2ARequestCallback is called before sending a request to the remote agent. +// +// If it returns non-nil result or error, the actual call is skipped and the returned value is used +// as the agent invocation result. +type BeforeA2ARequestCallback func(ctx agent.CallbackContext, req *a2a.SendMessageRequest) (*session.Event, error) + +// A2AEventConverter can be used to provide a custom implementation of A2A event transformation logic. +type A2AEventConverter func(ctx agent.InvocationContext, req *a2a.SendMessageRequest, event a2a.Event, err error) (*session.Event, error) + +// AfterA2ARequestCallback is called after receiving a response from the remote agent and converting it to a session.Event. +// In streaming responses the callback is invoked for every request. Session event parameter might be nil if conversion logic +// decides to not emit an A2A event. +// +// If it returns non-nil result or error, it gets emitted instead of the original result. +type AfterA2ARequestCallback func(ctx agent.CallbackContext, req *a2a.SendMessageRequest, resp *session.Event, err error) (*session.Event, error) + +// A2ARemoteTaskCleanupCallback is called if Run exited before a terminal event was received from the remote A2A server. +type A2ARemoteTaskCleanupCallback func(ctx context.Context, card *a2a.AgentCard, client A2AClient, taskInfo a2a.TaskInfo, cause error) + +// AgentCardProvider resolves an agent card on each agent invocation. +// Use [NewAgentCardProvider] to create a provider from a URL or file path. +// Callers that want lazy/cached resolution should implement caching within the provider function. +type AgentCardProvider func(ctx context.Context) (*a2a.AgentCard, error) + +// NewAgentCardProvider creates an [AgentCardProvider] that resolves an agent card from the given source. +// The source can be an http(s) URL or a local file path. +func NewAgentCardProvider(source string, opts ...agentcard.ResolveOption) AgentCardProvider { + return func(ctx context.Context) (*a2a.AgentCard, error) { + if strings.HasPrefix(source, "http://") || strings.HasPrefix(source, "https://") { + card, err := agentcard.DefaultResolver.Resolve(ctx, source, opts...) + if err != nil { + return nil, fmt.Errorf("failed to fetch an agent card: %w", err) + } + return card, nil + } + + fileBytes, err := os.ReadFile(source) + if err != nil { + return nil, fmt.Errorf("failed to read agent card from %q: %w", source, err) + } + + var card a2a.AgentCard + if err := json.Unmarshal(fileBytes, &card); err != nil { + return nil, fmt.Errorf("failed to unmarshal an agent card: %w", err) + } + return &card, nil + } +} + +// A2AConfig is used to describe and configure a remote agent. +type A2AConfig struct { + Name string + Description string + + // AgentCard is a static agent card. Either AgentCard or AgentCardProvider must be set. + AgentCard *a2a.AgentCard + // AgentCardProvider resolves an agent card on each agent invocation. + // Use [NewAgentCardProvider] to create a provider from a URL or file path. + // Either AgentCard or AgentCardProvider must be set. + AgentCardProvider AgentCardProvider + + // BeforeAgentCallbacks is a list of callbacks that are called sequentially + // before the agent starts its run. + // + // If any callback returns non-nil content or error, then the agent run and + // the remaining callbacks will be skipped, and a new event will be created + // from the content or error of that callback. + BeforeAgentCallbacks []agent.BeforeAgentCallback + // BeforeRequestCallbacks will be called in the order they are provided until + // there's a callback that returns a non-nil result or error. Then the + // actual request is skipped, and the returned response/error is used. + // + // This provides an opportunity to inspect, log, or modify the request object. + // It can also be used to implement caching by returning a cached + // response, which would skip the actual remote agent call. + BeforeRequestCallbacks []BeforeA2ARequestCallback + // Converter is used to convert a2a.Event to session.Event. If not provided, adka2a.ToSessionEvent + // is used as the default implementation and errors are converted to events with error payload. + Converter A2AEventConverter + // AfterRequestCallbacks will be called in the order they are provided until + // there's a callback that returns a non-nil result or error. Then + // the actual remote agent event is replaced with the returned result/error. + // + // This is the ideal place to log agent responses, collect metrics on token or perform + // pre-processing of events before a mapper is invoked. + AfterRequestCallbacks []AfterA2ARequestCallback + // AfterAgentCallbacks is a list of callbacks that are called sequentially + // after the agent has completed its run. + // + // If any callback returns non-nil content or error, then a new event will be + // created from the content or error of that callback and the remaining + // callbacks will be skipped. + AfterAgentCallbacks []agent.AfterAgentCallback + + // A2APartConverter is a custom converter for converting A2A parts to GenAI parts. + // Implementations should generally remember to leverage adka2a.ToGenAiPart for default conversions + // nil returns are considered intentionally dropped parts. + A2APartConverter adka2a.A2APartConverter + + // GenAIPartConverter is a custom converter for converting GenAI parts to A2A parts. + // Implementations should generally remember to leverage adka2a.ToA2APart for default conversions + // nil returns are considered intentionally dropped parts. + GenAIPartConverter adka2a.GenAIPartConverter + + // ClientProvider can be used to provide a custom implementation of A2A message sending. + ClientProvider A2AClientProvider + // MessageSendConfig is attached to a2a.SendMessageRequest sent on every agent invocation. + MessageSendConfig *a2a.SendMessageConfig + + // RemoteTaskCleanupCallback is called if Run exited before a terminal event was received from the remote A2A server. + // If Run exited due to an error including context cancellation it will be passed as cause. + // The context passed to this callback is the original context, but with Err() removed by context.WithoutCancel. + // If no callback is provided the default behavior is to make a cancel RPC request with 5 second timeout. + RemoteTaskCleanupCallback A2ARemoteTaskCleanupCallback +} + +// NewA2A creates a remote A2A agent. A2A (Agent-To-Agent) protocol is used for communication with an +// agent which can run in a different process or on a different host. +func NewA2A(cfg A2AConfig) (agent.Agent, error) { + if cfg.AgentCard == nil && cfg.AgentCardProvider == nil { + return nil, fmt.Errorf("either AgentCard or AgentCardProvider must be provided") + } + if cfg.ClientProvider == nil { + cfg.ClientProvider = NewA2AClientProvider(a2aclient.NewFactory()) + } + + remoteAgent := &a2aAgent{ + serverConfig: &iremoteagent.A2AServerConfig{ + AgentCard: cfg.AgentCard, + AgentCardProvider: cfg.AgentCardProvider, + ClientProvider: cfg.ClientProvider, + }, + } + agent, err := agent.New(agent.Config{ + Name: cfg.Name, + Description: cfg.Description, + BeforeAgentCallbacks: cfg.BeforeAgentCallbacks, + AfterAgentCallbacks: cfg.AfterAgentCallbacks, + Run: func(ic agent.InvocationContext) iter.Seq2[*session.Event, error] { + return remoteAgent.run(ic, cfg) + }, + }) + if err != nil { + return nil, err + } + + internalAgent, ok := agent.(agentinternal.Agent) + if !ok { + return nil, fmt.Errorf("internal error: failed to convert to internal agent") + } + state := agentinternal.Reveal(internalAgent) + state.AgentType = agentinternal.TypeRemoteAgent + state.Config = iremoteagent.RemoteAgentState{A2A: remoteAgent.serverConfig} + + return agent, nil +} + +type a2aAgent struct { + serverConfig *iremoteagent.A2AServerConfig +} + +func (a *a2aAgent) run(ctx agent.InvocationContext, cfg A2AConfig) iter.Seq2[*session.Event, error] { + return func(yield func(*session.Event, error) bool) { + card, err := iremoteagent.ResolveAgentCard(ctx, a.serverConfig) + if err != nil { + yield(toErrorEvent(ctx, fmt.Errorf("agent card resolution failed: %w", err)), nil) + return + } + + sender, err := cfg.ClientProvider(ctx, card) + if err != nil { + yield(toErrorEvent(ctx, fmt.Errorf("sender creation failed: %w", err)), nil) + return + } + defer destroy(ctx, sender) + + msg, err := newMessage(ctx, cfg) + if err != nil { + yield(toErrorEvent(ctx, fmt.Errorf("message creation failed: %w", err)), nil) + return + } + + req := &a2a.SendMessageRequest{Message: msg, Config: cfg.MessageSendConfig} + processor := newRunProcessor(cfg, req) + + if bcbResp, bcbErr := processor.runBeforeA2ARequestCallbacks(ctx); bcbResp != nil || bcbErr != nil { + if acbResp, acbErr := processor.runAfterA2ARequestCallbacks(ctx, bcbResp, bcbErr); acbResp != nil || acbErr != nil { + yield(acbResp, acbErr) + } else { + yield(bcbResp, bcbErr) + } + return + } + + if len(msg.Parts) == 0 { + resp := adka2a.NewRemoteAgentEvent(ctx) + if cbResp, cbErr := processor.runAfterA2ARequestCallbacks(ctx, resp, err); cbResp != nil || cbErr != nil { + yield(cbResp, cbErr) + } else { + yield(resp, nil) + } + return + } + + var lastErr error + yieldErr := func(err error) bool { + lastErr = err + return yield(nil, err) + } + + var lastEvent a2a.Event + defer func() { + err := lastErr + if err == nil && ctx.Err() != nil { + err = context.Cause(ctx) + } + cleanupRemoteTask(ctx, cfg, card, sender, lastEvent, err) + }() + + processEvent := func(a2aEvent a2a.Event, a2aErr error) bool { + if a2aEvent != nil { + lastEvent = a2aEvent + } + + var err error + var event *session.Event + if cfg.Converter != nil { + event, err = cfg.Converter(ctx, req, a2aEvent, a2aErr) + } else { + event, err = processor.convertToSessionEvent(ctx, a2aEvent, a2aErr) + } + + if cbResp, cbErr := processor.runAfterA2ARequestCallbacks(ctx, event, err); cbResp != nil || cbErr != nil { + if cbErr != nil { + return yieldErr(cbErr) + } + event = cbResp + err = nil + } + + if err != nil { + return yieldErr(err) + } + + if event != nil { // an event might be skipped + for _, toEmit := range processor.aggregatePartial(ctx, a2aEvent, event) { + if !yield(toEmit, nil) { + return false + } + } + } + return true + } + + if ctx.RunConfig().StreamingMode == agent.StreamingModeNone { + a2aEvent, a2aErr := sender.SendMessage(ctx, req) + processEvent(a2aEvent, a2aErr) + return + } + + for a2aEvent, a2aErr := range sender.SendStreamingMessage(ctx, req) { + if !processEvent(a2aEvent, a2aErr) { + return + } + } + } +} + +func cleanupRemoteTask(ctx context.Context, cfg A2AConfig, card *a2a.AgentCard, client A2AClient, lastEvent a2a.Event, cause error) { + if lastEvent == nil { + return + } + taskID := lastEvent.TaskInfo().TaskID + if taskID == "" { + return + } + if _, ok := lastEvent.(*a2a.Message); ok { + return + } + var state a2a.TaskState + if tu, ok := lastEvent.(*a2a.TaskStatusUpdateEvent); ok { + state = tu.Status.State + } + if t, ok := lastEvent.(*a2a.Task); ok { + state = t.Status.State + } + if state.Terminal() { + return + } + + ctx = context.WithoutCancel(ctx) + + if cfg.RemoteTaskCleanupCallback != nil { + cfg.RemoteTaskCleanupCallback(ctx, card, client, lastEvent.TaskInfo(), cause) + return + } + + if state == a2a.TaskStateInputRequired && cause == nil { + return + } + cancelCtx, cancelTimeout := context.WithTimeout(ctx, 5*time.Second) + defer cancelTimeout() + _, err := client.CancelTask(cancelCtx, &a2a.CancelTaskRequest{ID: taskID}) + if err != nil { + log.Warn(ctx, "failed to cancel task", "task_id", taskID, "error", err) + } +} + +func newMessage(ctx agent.InvocationContext, cfg A2AConfig) (*a2a.Message, error) { + events := ctx.Session().Events() + if userFnCall := getUserFunctionCallAt(events, events.Len()-1); userFnCall != nil { + event := userFnCall.response + parts, err := convertParts(ctx, cfg, event) + if err != nil { + return nil, fmt.Errorf("event part conversion failed: %w", err) + } + msg := a2a.NewMessage(a2a.MessageRoleUser, parts...) + msg.TaskID = a2a.TaskID(userFnCall.taskID) + msg.ContextID = userFnCall.contextID + return msg, nil + } + + parts, contextID := toMissingRemoteSessionParts(ctx, events, cfg) + msg := a2a.NewMessage(a2a.MessageRoleUser, parts...) + msg.ContextID = contextID + return msg, nil +} + +func toErrorEvent(ctx agent.InvocationContext, err error) *session.Event { + event := adka2a.NewRemoteAgentEvent(ctx) + event.ErrorMessage = err.Error() + event.CustomMetadata = map[string]any{adka2a.ToADKMetaKey("error"): err.Error()} + event.TurnComplete = true + return event +} + +func destroy(ctx context.Context, client A2AClient) { + if err := client.Destroy(); err != nil { + log.Warn(ctx, "failed to destroy client", "error", err) + } +} diff --git a/agent/remoteagent/a2a_agent_run_processor.go b/agent/remoteagent/v2/a2a_agent_run_processor.go similarity index 96% rename from agent/remoteagent/a2a_agent_run_processor.go rename to agent/remoteagent/v2/a2a_agent_run_processor.go index 615a3d1ea..dc3b58faa 100644 --- a/agent/remoteagent/a2a_agent_run_processor.go +++ b/agent/remoteagent/v2/a2a_agent_run_processor.go @@ -19,13 +19,13 @@ import ( "maps" "slices" - "github.com/a2aproject/a2a-go/a2a" + "github.com/a2aproject/a2a-go/v2/a2a" "google.golang.org/genai" "google.golang.org/adk/agent" icontext "google.golang.org/adk/internal/context" "google.golang.org/adk/internal/converters" - "google.golang.org/adk/server/adka2a" + "google.golang.org/adk/server/adka2a/v2" "google.golang.org/adk/session" ) @@ -40,7 +40,7 @@ type artifactAggregation struct { type a2aAgentRunProcessor struct { config A2AConfig partConverter adka2a.A2APartConverter - request *a2a.MessageSendParams + request *a2a.SendMessageRequest // partial event contents emitted before the terminal event aggregations map[a2a.ArtifactID]*artifactAggregation @@ -48,7 +48,7 @@ type a2aAgentRunProcessor struct { aggregationOrder []a2a.ArtifactID } -func newRunProcessor(config A2AConfig, request *a2a.MessageSendParams) *a2aAgentRunProcessor { +func newRunProcessor(config A2AConfig, request *a2a.SendMessageRequest) *a2aAgentRunProcessor { return &a2aAgentRunProcessor{ config: config, request: request, @@ -68,7 +68,7 @@ func (p *a2aAgentRunProcessor) aggregatePartial(ctx agent.InvocationContext, a2a } // RemoteAgent event stream finished, emit any aggregated events data we have before the final event - if statusUpdate, ok := a2aEvent.(*a2a.TaskStatusUpdateEvent); ok && statusUpdate.Final { + if statusUpdate, ok := a2aEvent.(*a2a.TaskStatusUpdateEvent); ok && statusUpdate.Status.State.Terminal() { var events []*session.Event for _, aid := range p.aggregationOrder { if agg, ok := p.aggregations[aid]; ok { @@ -216,6 +216,9 @@ func (p *a2aAgentRunProcessor) runBeforeA2ARequestCallbacks(ctx agent.Invocation } func (p *a2aAgentRunProcessor) runAfterA2ARequestCallbacks(ctx agent.InvocationContext, resp *session.Event, err error) (*session.Event, error) { + if resp == nil && err == nil { + return nil, nil + } cctx := icontext.NewCallbackContext(ctx) for _, callback := range p.config.AfterRequestCallbacks { if cbEvent, cbErr := callback(cctx, p.request, resp, err); cbEvent != nil || cbErr != nil { diff --git a/agent/remoteagent/a2a_agent_run_processor_test.go b/agent/remoteagent/v2/a2a_agent_run_processor_test.go similarity index 88% rename from agent/remoteagent/a2a_agent_run_processor_test.go rename to agent/remoteagent/v2/a2a_agent_run_processor_test.go index 03f598d97..2e6351c78 100644 --- a/agent/remoteagent/a2a_agent_run_processor_test.go +++ b/agent/remoteagent/v2/a2a_agent_run_processor_test.go @@ -17,7 +17,7 @@ package remoteagent import ( "testing" - "github.com/a2aproject/a2a-go/a2a" + "github.com/a2aproject/a2a-go/v2/a2a" "github.com/google/go-cmp/cmp" "github.com/google/go-cmp/cmp/cmpopts" "google.golang.org/genai" @@ -26,7 +26,7 @@ import ( icontext "google.golang.org/adk/internal/context" "google.golang.org/adk/internal/utils" "google.golang.org/adk/model" - "google.golang.org/adk/server/adka2a" + "google.golang.org/adk/server/adka2a/v2" "google.golang.org/adk/session" ) @@ -40,7 +40,7 @@ func TestA2AAgentRunProcessor_aggregatePartial(t *testing.T) { return &a2a.TaskArtifactUpdateEvent{ TaskID: task.ID, ContextID: task.ContextID, - Artifact: &a2a.Artifact{ID: aid, Parts: a2a.ContentParts{a2a.TextPart{Text: text}}}, + Artifact: &a2a.Artifact{ID: aid, Parts: a2a.ContentParts{a2a.NewTextPart(text)}}, LastChunk: flags.lastChunk, Append: flags.append, } @@ -80,9 +80,9 @@ func TestA2AAgentRunProcessor_aggregatePartial(t *testing.T) { { name: "do not aggregate when ADK events", events: []a2a.Event{ - withADKPartial(a2a.NewArtifactUpdateEvent(task, aid1, a2a.TextPart{Text: "Hel"}), true), - withADKPartial(a2a.NewArtifactUpdateEvent(task, aid1, a2a.TextPart{Text: "lo"}), true), - withADKPartial(a2a.NewArtifactUpdateEvent(task, aid1, a2a.TextPart{Text: "Hello"}), false), + withADKPartial(a2a.NewArtifactUpdateEvent(task, aid1, a2a.NewTextPart("Hel")), true), + withADKPartial(a2a.NewArtifactUpdateEvent(task, aid1, a2a.NewTextPart("lo")), true), + withADKPartial(a2a.NewArtifactUpdateEvent(task, aid1, a2a.NewTextPart("Hello")), false), newFinalStatusUpdate(task, a2a.TaskStateCompleted), }, wantEvents: []*session.Event{ @@ -95,10 +95,10 @@ func TestA2AAgentRunProcessor_aggregatePartial(t *testing.T) { { name: "aggregation reset by final snapshot", events: []a2a.Event{ - a2a.NewArtifactUpdateEvent(task, "a1", a2a.TextPart{Text: "ignore me"}), + a2a.NewArtifactUpdateEvent(task, "a1", a2a.NewTextPart("ignore me")), &a2a.Task{ ID: task.ID, - Artifacts: []*a2a.Artifact{{Parts: a2a.ContentParts{a2a.TextPart{Text: "done"}}}}, + Artifacts: []*a2a.Artifact{{Parts: a2a.ContentParts{a2a.NewTextPart("done")}}}, Status: a2a.TaskStatus{State: a2a.TaskStateCompleted}, }, }, @@ -110,9 +110,9 @@ func TestA2AAgentRunProcessor_aggregatePartial(t *testing.T) { { name: "aggregation reset by non-final snapshot", events: []a2a.Event{ - a2a.NewArtifactUpdateEvent(task, "a1", a2a.TextPart{Text: "foo"}), + a2a.NewArtifactUpdateEvent(task, "a1", a2a.NewTextPart("foo")), &a2a.Task{ID: task.ID}, - a2a.NewArtifactUpdateEvent(task, "a1", a2a.TextPart{Text: "bar"}), + a2a.NewArtifactUpdateEvent(task, "a1", a2a.NewTextPart("bar")), newFinalStatusUpdate(task, a2a.TaskStateCompleted), }, wantEvents: []*session.Event{ @@ -232,11 +232,12 @@ func TestA2AAgentRunProcessor_aggregatePartial(t *testing.T) { { name: "thoughts aggregation", events: []a2a.Event{ - a2a.NewArtifactUpdateEvent(task, "a1", a2a.TextPart{ - Text: "thinking...", - Metadata: map[string]any{adka2a.ToA2AMetaKey("thought"): true}, - }), - a2a.NewArtifactUpdateEvent(task, "a1", a2a.TextPart{Text: "done"}), + func() *a2a.TaskArtifactUpdateEvent { + p := a2a.NewTextPart("thinking...") + p.SetMeta(adka2a.ToA2AMetaKey("thought"), true) + return a2a.NewArtifactUpdateEvent(task, "a1", p) + }(), + a2a.NewArtifactUpdateEvent(task, "a1", a2a.NewTextPart("done")), newFinalStatusUpdate(task, a2a.TaskStateCompleted), }, wantEvents: []*session.Event{ @@ -255,16 +256,16 @@ func TestA2AAgentRunProcessor_aggregatePartial(t *testing.T) { { name: "interleaved thought and text", events: []a2a.Event{ - a2a.NewArtifactUpdateEvent(task, "a1", a2a.TextPart{ - Text: "thinking1", + a2a.NewArtifactUpdateEvent(task, "a1", &a2a.Part{ + Content: a2a.Text("thinking1"), Metadata: map[string]any{adka2a.ToA2AMetaKey("thought"): true}, }), - a2a.NewArtifactUpdateEvent(task, "a1", a2a.TextPart{Text: "text1"}), - a2a.NewArtifactUpdateEvent(task, "a1", a2a.TextPart{ - Text: "thinking2", + a2a.NewArtifactUpdateEvent(task, "a1", a2a.NewTextPart("text1")), + a2a.NewArtifactUpdateEvent(task, "a1", &a2a.Part{ + Content: a2a.Text("thinking2"), Metadata: map[string]any{adka2a.ToA2AMetaKey("thought"): true}, }), - a2a.NewArtifactUpdateEvent(task, "a1", a2a.TextPart{Text: "text2"}), + a2a.NewArtifactUpdateEvent(task, "a1", a2a.NewTextPart("text2")), newFinalStatusUpdate(task, a2a.TaskStateCompleted), }, wantEvents: []*session.Event{ diff --git a/agent/remoteagent/a2a_agent_test.go b/agent/remoteagent/v2/a2a_agent_test.go similarity index 78% rename from agent/remoteagent/a2a_agent_test.go rename to agent/remoteagent/v2/a2a_agent_test.go index a3f484538..fabed995c 100644 --- a/agent/remoteagent/a2a_agent_test.go +++ b/agent/remoteagent/v2/a2a_agent_test.go @@ -26,10 +26,9 @@ import ( "testing" "time" - "github.com/a2aproject/a2a-go/a2a" - "github.com/a2aproject/a2a-go/a2aclient" - "github.com/a2aproject/a2a-go/a2asrv" - "github.com/a2aproject/a2a-go/a2asrv/eventqueue" + "github.com/a2aproject/a2a-go/v2/a2a" + "github.com/a2aproject/a2a-go/v2/a2aclient" + "github.com/a2aproject/a2a-go/v2/a2asrv" "github.com/google/go-cmp/cmp" "github.com/google/go-cmp/cmp/cmpopts" "google.golang.org/genai" @@ -39,35 +38,37 @@ import ( "google.golang.org/adk/internal/utils" "google.golang.org/adk/model" "google.golang.org/adk/runner" - "google.golang.org/adk/server/adka2a" + "google.golang.org/adk/server/adka2a/v2" "google.golang.org/adk/session" ) type mockA2AExecutor struct { - executeFn func(ctx context.Context, reqCtx *a2asrv.RequestContext, queue eventqueue.Queue) error - cancelFn func(ctx context.Context, reqCtx *a2asrv.RequestContext, queue eventqueue.Queue) error - cleanupFn func(ctx context.Context, reqCtx *a2asrv.RequestContext, result a2a.SendMessageResult, cause error) + executeFn func(ctx context.Context, execCtx *a2asrv.ExecutorContext) iter.Seq2[a2a.Event, error] + cancelFn func(ctx context.Context, reqCtx *a2asrv.ExecutorContext) iter.Seq2[a2a.Event, error] + cleanupFn func(ctx context.Context, reqCtx *a2asrv.ExecutorContext, result a2a.SendMessageResult, cause error) } var _ a2asrv.AgentExecutor = (*mockA2AExecutor)(nil) -func (e *mockA2AExecutor) Execute(ctx context.Context, reqCtx *a2asrv.RequestContext, queue eventqueue.Queue) error { +func (e *mockA2AExecutor) Execute(ctx context.Context, execCtx *a2asrv.ExecutorContext) iter.Seq2[a2a.Event, error] { if e.executeFn != nil { - return e.executeFn(ctx, reqCtx, queue) + return e.executeFn(ctx, execCtx) + } + return func(yield func(a2a.Event, error) bool) { + yield(nil, fmt.Errorf("not implemented")) } - return fmt.Errorf("not implemented") } -func (e *mockA2AExecutor) Cancel(ctx context.Context, reqCtx *a2asrv.RequestContext, queue eventqueue.Queue) error { +func (e *mockA2AExecutor) Cancel(ctx context.Context, execCtx *a2asrv.ExecutorContext) iter.Seq2[a2a.Event, error] { if e.cancelFn != nil { - return e.cancelFn(ctx, reqCtx, queue) + return e.cancelFn(ctx, execCtx) + } + return func(yield func(a2a.Event, error) bool) { + yield(a2a.NewStatusUpdateEvent(execCtx, a2a.TaskStateCanceled, nil), fmt.Errorf("not implemented")) } - ev := a2a.NewStatusUpdateEvent(reqCtx, a2a.TaskStateCanceled, nil) - ev.Final = true - return queue.Write(ctx, ev) } -func (e *mockA2AExecutor) Cleanup(ctx context.Context, reqCtx *a2asrv.RequestContext, result a2a.SendMessageResult, cause error) { +func (e *mockA2AExecutor) Cleanup(ctx context.Context, reqCtx *a2asrv.ExecutorContext, result a2a.SendMessageResult, cause error) { if e.cleanupFn != nil { e.cleanupFn(ctx, reqCtx, result, cause) } @@ -88,7 +89,7 @@ func startA2AServer(agentExecutor a2asrv.AgentExecutor) *testA2AServer { func newA2ARemoteAgent(t *testing.T, name string, server *testA2AServer) agent.Agent { t.Helper() - card := &a2a.AgentCard{PreferredTransport: a2a.TransportProtocolJSONRPC, URL: server.URL, Capabilities: a2a.AgentCapabilities{Streaming: true}} + card := &a2a.AgentCard{SupportedInterfaces: []*a2a.AgentInterface{a2a.NewAgentInterface(server.URL, a2a.TransportProtocolJSONRPC)}, Capabilities: a2a.AgentCapabilities{Streaming: true}} return utils.Must(NewA2A(A2AConfig{AgentCard: card, Name: name})) } @@ -168,45 +169,54 @@ func newADKEventReplay(t *testing.T, name string, events []*session.Event) agent func newA2AEventReplay(t *testing.T, events []a2a.Event) a2asrv.AgentExecutor { return &mockA2AExecutor{ - executeFn: func(ctx context.Context, reqCtx *a2asrv.RequestContext, queue eventqueue.Queue) error { - for _, ev := range events { - // A2A stack is going to fail the request if events don't have correct taskID and contextID - switch v := ev.(type) { - case *a2a.Message: - v.TaskID = reqCtx.TaskID - v.ContextID = reqCtx.ContextID - case *a2a.Task: - v.ID = reqCtx.TaskID - v.ContextID = reqCtx.ContextID - case *a2a.TaskStatusUpdateEvent: - v.TaskID = reqCtx.TaskID - v.ContextID = reqCtx.ContextID - case *a2a.TaskArtifactUpdateEvent: - v.TaskID = reqCtx.TaskID - v.ContextID = reqCtx.ContextID + executeFn: func(ctx context.Context, execCtx *a2asrv.ExecutorContext) iter.Seq2[a2a.Event, error] { + return func(yield func(a2a.Event, error) bool) { + if len(events) > 0 { + if _, ok := events[0].(*a2a.Task); !ok { + if _, ok := events[0].(*a2a.Message); !ok { + if !yield(a2a.NewSubmittedTask(execCtx, execCtx.Message), nil) { + return + } + } + } } - if err := queue.Write(ctx, ev); err != nil { - t.Errorf("queue.Write() error = %v", err) + for _, ev := range events { + // A2A stack is going to fail the request if events don't have correct taskID and contextID + switch v := ev.(type) { + case *a2a.Message: + v.TaskID = execCtx.TaskID + v.ContextID = execCtx.ContextID + case *a2a.Task: + v.ID = execCtx.TaskID + v.ContextID = execCtx.ContextID + case *a2a.TaskStatusUpdateEvent: + v.TaskID = execCtx.TaskID + v.ContextID = execCtx.ContextID + case *a2a.TaskArtifactUpdateEvent: + v.TaskID = execCtx.TaskID + v.ContextID = execCtx.ContextID + } + if !yield(ev, nil) { + return + } } } - return nil }, } } func newUserHello() *session.Event { - event := session.NewEvent("invocation") + event := session.NewEventWithContext(context.Background(), "invocation") event.Author = "user" event.Content = genai.NewContentFromText("hello", genai.RoleUser) return event } -func newFinalStatusUpdate(task *a2a.Task, state a2a.TaskState, msgParts ...a2a.Part) *a2a.TaskStatusUpdateEvent { +func newFinalStatusUpdate(task *a2a.Task, state a2a.TaskState, msgParts ...*a2a.Part) *a2a.TaskStatusUpdateEvent { event := a2a.NewStatusUpdateEvent(task, state, nil) if len(msgParts) > 0 { event.Status.Message = a2a.NewMessageForTask(a2a.MessageRoleAgent, task, msgParts...) } - event.Final = true return event } @@ -430,6 +440,11 @@ func TestRemoteAgent_ADK2ADK(t *testing.T) { func TestRemoteAgent_ADK2A2A(t *testing.T) { task := &a2a.Task{ID: a2a.NewTaskID(), ContextID: a2a.NewContextID()} artifactEvent := a2a.NewArtifactEvent(task) + newArtifactEvent := func(parts ...*a2a.Part) *a2a.TaskArtifactUpdateEvent { + event := a2a.NewArtifactUpdateEvent(task, artifactEvent.Artifact.ID, parts...) + event.Append = false + return event + } testCases := []struct { name string @@ -444,7 +459,7 @@ func TestRemoteAgent_ADK2A2A(t *testing.T) { { name: "message", remoteEvents: []a2a.Event{ - a2a.NewMessage(a2a.MessageRoleAgent, a2a.TextPart{Text: "hello"}, a2a.TextPart{Text: "world"}), + a2a.NewMessage(a2a.MessageRoleAgent, a2a.NewTextPart("hello"), a2a.NewTextPart("world")), }, wantResponses: []model.LLMResponse{ { @@ -468,7 +483,7 @@ func TestRemoteAgent_ADK2A2A(t *testing.T) { remoteEvents: []a2a.Event{ &a2a.Task{Status: a2a.TaskStatus{ State: a2a.TaskStateCompleted, - Message: a2a.NewMessage(a2a.MessageRoleAgent, a2a.TextPart{Text: "hello"}), + Message: a2a.NewMessage(a2a.MessageRoleAgent, a2a.NewTextPart("hello")), }}, }, wantResponses: []model.LLMResponse{{ @@ -482,7 +497,7 @@ func TestRemoteAgent_ADK2A2A(t *testing.T) { &a2a.Task{ Status: a2a.TaskStatus{State: a2a.TaskStateCompleted}, Artifacts: []*a2a.Artifact{ - {Parts: a2a.ContentParts{a2a.TextPart{Text: "hello"}, a2a.TextPart{Text: "world"}}}, + {Parts: a2a.ContentParts{a2a.NewTextPart("hello"), a2a.NewTextPart("world")}}, }, }, }, @@ -501,12 +516,12 @@ func TestRemoteAgent_ADK2A2A(t *testing.T) { remoteEvents: []a2a.Event{ &a2a.Task{Status: a2a.TaskStatus{ State: a2a.TaskStateWorking, - Message: a2a.NewMessage(a2a.MessageRoleAgent, a2a.TextPart{Text: "hello"}), + Message: a2a.NewMessage(a2a.MessageRoleAgent, a2a.NewTextPart("hello")), }}, &a2a.Task{ Status: a2a.TaskStatus{State: a2a.TaskStateCompleted}, Artifacts: []*a2a.Artifact{ - {Parts: a2a.ContentParts{a2a.TextPart{Text: "world"}}}, + {Parts: a2a.ContentParts{a2a.NewTextPart("world")}}, }, }, }, @@ -521,8 +536,8 @@ func TestRemoteAgent_ADK2A2A(t *testing.T) { &a2a.Task{ Status: a2a.TaskStatus{State: a2a.TaskStateCompleted}, Artifacts: []*a2a.Artifact{ - {Parts: a2a.ContentParts{a2a.TextPart{Text: "hello"}}}, - {Parts: a2a.ContentParts{a2a.TextPart{Text: "world"}}}, + {Parts: a2a.ContentParts{a2a.NewTextPart("hello")}}, + {Parts: a2a.ContentParts{a2a.NewTextPart("world")}}, }, }, }, @@ -539,9 +554,8 @@ func TestRemoteAgent_ADK2A2A(t *testing.T) { { name: "artifact parts translation", remoteEvents: []a2a.Event{ - artifactEvent, - a2a.NewArtifactUpdateEvent(task, artifactEvent.Artifact.ID, a2a.TextPart{Text: "hello"}), - a2a.NewArtifactUpdateEvent(task, artifactEvent.Artifact.ID, a2a.TextPart{Text: "world"}), + newArtifactEvent(a2a.NewTextPart("hello")), + a2a.NewArtifactUpdateEvent(task, artifactEvent.Artifact.ID, a2a.NewTextPart("world")), newFinalStatusUpdate(task, a2a.TaskStateCompleted), }, wantResponses: []model.LLMResponse{ @@ -554,9 +568,9 @@ func TestRemoteAgent_ADK2A2A(t *testing.T) { { name: "non-final status update messages as thoughts", remoteEvents: []a2a.Event{ - a2a.NewStatusUpdateEvent(task, a2a.TaskStateSubmitted, a2a.NewMessage(a2a.MessageRoleAgent, a2a.TextPart{Text: "submitted...\n"})), - a2a.NewStatusUpdateEvent(task, a2a.TaskStateWorking, a2a.NewMessage(a2a.MessageRoleAgent, a2a.TextPart{Text: "working...\n"})), - newFinalStatusUpdate(task, a2a.TaskStateCompleted, a2a.TextPart{Text: "completed!"}), + a2a.NewStatusUpdateEvent(task, a2a.TaskStateSubmitted, a2a.NewMessage(a2a.MessageRoleAgent, a2a.NewTextPart("submitted...\n"))), + a2a.NewStatusUpdateEvent(task, a2a.TaskStateWorking, a2a.NewMessage(a2a.MessageRoleAgent, a2a.NewTextPart("working...\n"))), + newFinalStatusUpdate(task, a2a.TaskStateCompleted, a2a.NewTextPart("completed!")), }, wantResponses: []model.LLMResponse{ {Content: &genai.Content{Parts: []*genai.Part{{Text: "submitted...\n", Thought: true}}, Role: genai.RoleModel}, Partial: true}, @@ -578,32 +592,27 @@ func TestRemoteAgent_ADK2A2A(t *testing.T) { { name: "partial and non-partial event aggregation", remoteEvents: []a2a.Event{ - artifactEvent, - &a2a.TaskArtifactUpdateEvent{ - TaskID: task.ID, - ContextID: task.ContextID, - Artifact: &a2a.Artifact{ID: artifactEvent.Artifact.ID, Parts: a2a.ContentParts{a2a.TextPart{Text: "1"}}}, - Append: true, - }, + newArtifactEvent(a2a.NewTextPart("1")), &a2a.TaskArtifactUpdateEvent{ TaskID: task.ID, ContextID: task.ContextID, - Artifact: &a2a.Artifact{ID: artifactEvent.Artifact.ID, Parts: a2a.ContentParts{a2a.TextPart{Text: "2"}}}, + Artifact: &a2a.Artifact{ID: artifactEvent.Artifact.ID, Parts: a2a.ContentParts{a2a.NewTextPart("2")}}, Append: true, }, + &a2a.TaskArtifactUpdateEvent{ TaskID: task.ID, ContextID: task.ContextID, - Artifact: &a2a.Artifact{ID: artifactEvent.Artifact.ID, Parts: a2a.ContentParts{a2a.TextPart{Text: "3"}}}, + Artifact: &a2a.Artifact{ID: artifactEvent.Artifact.ID, Parts: a2a.ContentParts{a2a.NewTextPart("3")}}, Append: false, }, &a2a.TaskArtifactUpdateEvent{ TaskID: task.ID, ContextID: task.ContextID, - Artifact: &a2a.Artifact{ID: artifactEvent.Artifact.ID, Parts: a2a.ContentParts{a2a.TextPart{Text: "4"}}}, + Artifact: &a2a.Artifact{ID: artifactEvent.Artifact.ID, Parts: a2a.ContentParts{a2a.NewTextPart("4")}}, Append: true, }, - newFinalStatusUpdate(task, a2a.TaskStateCompleted, a2a.TextPart{Text: "5"}), + newFinalStatusUpdate(task, a2a.TaskStateCompleted, a2a.NewTextPart("5")), }, wantResponses: []model.LLMResponse{ {Content: genai.NewContentFromText("1", genai.RoleModel), Partial: true}, @@ -654,7 +663,7 @@ func TestRemoteAgent_RequestCallbacks(t *testing.T) { testCases := []struct { name string sessionEvents []*session.Event - events func(*a2asrv.RequestContext) []a2a.Event + events func(*a2asrv.ExecutorContext) []a2a.Event before []BeforeA2ARequestCallback after []AfterA2ARequestCallback converter A2AEventConverter @@ -663,17 +672,17 @@ func TestRemoteAgent_RequestCallbacks(t *testing.T) { }{ { name: "request and response modification", - events: func(rc *a2asrv.RequestContext) []a2a.Event { - return []a2a.Event{a2a.NewMessage(a2a.MessageRoleAgent, a2a.TextPart{Text: "foo"})} + events: func(rc *a2asrv.ExecutorContext) []a2a.Event { + return []a2a.Event{a2a.NewMessage(a2a.MessageRoleAgent, a2a.NewTextPart("foo"))} }, before: []BeforeA2ARequestCallback{ - func(ctx agent.CallbackContext, req *a2a.MessageSendParams) (*session.Event, error) { + func(ctx agent.CallbackContext, req *a2a.SendMessageRequest) (*session.Event, error) { req.Metadata = map[string]any{"counter": 1} return nil, nil }, }, after: []AfterA2ARequestCallback{ - func(ctx agent.CallbackContext, req *a2a.MessageSendParams, result *session.Event, err error) (*session.Event, error) { + func(ctx agent.CallbackContext, req *a2a.SendMessageRequest, result *session.Event, err error) (*session.Event, error) { result.Content = genai.NewContentFromText(result.Content.Parts[0].Text+"bar", genai.RoleModel) result.CustomMetadata = req.Metadata return nil, nil @@ -689,18 +698,17 @@ func TestRemoteAgent_RequestCallbacks(t *testing.T) { }, { name: "after invoked for every event", - events: func(rc *a2asrv.RequestContext) []a2a.Event { - artifactEvent := a2a.NewArtifactEvent(rc, a2a.TextPart{Text: "Hello"}) - finalEvent := a2a.NewStatusUpdateEvent(rc, a2a.TaskStateCompleted, nil) - finalEvent.Final = true + events: func(rc *a2asrv.ExecutorContext) []a2a.Event { + artifactEvent := a2a.NewArtifactEvent(rc, a2a.NewTextPart("Hello")) return []a2a.Event{ + a2a.NewSubmittedTask(rc, a2a.NewMessage(a2a.MessageRoleUser, a2a.NewTextPart("..."))), artifactEvent, - a2a.NewArtifactUpdateEvent(rc, artifactEvent.Artifact.ID, a2a.TextPart{Text: ", world!"}), - finalEvent, + a2a.NewArtifactUpdateEvent(rc, artifactEvent.Artifact.ID, a2a.NewTextPart(", world!")), + a2a.NewStatusUpdateEvent(rc, a2a.TaskStateCompleted, nil), } }, after: []AfterA2ARequestCallback{ - func(ctx agent.CallbackContext, req *a2a.MessageSendParams, result *session.Event, err error) (*session.Event, error) { + func(ctx agent.CallbackContext, req *a2a.SendMessageRequest, result *session.Event, err error) (*session.Event, error) { result.CustomMetadata = map[string]any{"foo": "bar"} return nil, nil }, @@ -728,16 +736,15 @@ func TestRemoteAgent_RequestCallbacks(t *testing.T) { }, { name: "after error stops the run", - events: func(rc *a2asrv.RequestContext) []a2a.Event { + events: func(rc *a2asrv.ExecutorContext) []a2a.Event { finalEvent := a2a.NewStatusUpdateEvent(rc, a2a.TaskStateCompleted, nil) - finalEvent.Final = true return []a2a.Event{ - a2a.NewArtifactEvent(rc, a2a.TextPart{Text: "Hello"}), + a2a.NewArtifactEvent(rc, a2a.NewTextPart("Hello")), finalEvent, } }, after: []AfterA2ARequestCallback{ - func(ctx agent.CallbackContext, req *a2a.MessageSendParams, result *session.Event, err error) (*session.Event, error) { + func(ctx agent.CallbackContext, req *a2a.SendMessageRequest, result *session.Event, err error) (*session.Event, error) { return nil, fmt.Errorf("rejected") }, }, @@ -746,7 +753,7 @@ func TestRemoteAgent_RequestCallbacks(t *testing.T) { { name: "request overwrite with response", before: []BeforeA2ARequestCallback{ - func(ctx agent.CallbackContext, req *a2a.MessageSendParams) (*session.Event, error) { + func(ctx agent.CallbackContext, req *a2a.SendMessageRequest) (*session.Event, error) { return &session.Event{LLMResponse: model.LLMResponse{Content: genai.NewContentFromText("hello", genai.RoleModel)}}, nil }, }, @@ -755,7 +762,7 @@ func TestRemoteAgent_RequestCallbacks(t *testing.T) { { name: "request overwrite with error", before: []BeforeA2ARequestCallback{ - func(ctx agent.CallbackContext, req *a2a.MessageSendParams) (*session.Event, error) { + func(ctx agent.CallbackContext, req *a2a.SendMessageRequest) (*session.Event, error) { return nil, fmt.Errorf("failed") }, }, @@ -764,7 +771,7 @@ func TestRemoteAgent_RequestCallbacks(t *testing.T) { { name: "response overwrite", after: []AfterA2ARequestCallback{ - func(ctx agent.CallbackContext, req *a2a.MessageSendParams, result *session.Event, err error) (*session.Event, error) { + func(ctx agent.CallbackContext, req *a2a.SendMessageRequest, result *session.Event, err error) (*session.Event, error) { return &session.Event{LLMResponse: model.LLMResponse{Content: genai.NewContentFromText("hello", genai.RoleModel)}}, nil }, }, @@ -773,7 +780,7 @@ func TestRemoteAgent_RequestCallbacks(t *testing.T) { { name: "response overwrite with error", after: []AfterA2ARequestCallback{ - func(ctx agent.CallbackContext, req *a2a.MessageSendParams, result *session.Event, err error) (*session.Event, error) { + func(ctx agent.CallbackContext, req *a2a.SendMessageRequest, result *session.Event, err error) (*session.Event, error) { return nil, fmt.Errorf("failed") }, }, @@ -782,10 +789,10 @@ func TestRemoteAgent_RequestCallbacks(t *testing.T) { { name: "before interceptor short-circuit", before: []BeforeA2ARequestCallback{ - func(ctx agent.CallbackContext, req *a2a.MessageSendParams) (*session.Event, error) { + func(ctx agent.CallbackContext, req *a2a.SendMessageRequest) (*session.Event, error) { return nil, fmt.Errorf("failed") }, - func(ctx agent.CallbackContext, req *a2a.MessageSendParams) (*session.Event, error) { + func(ctx agent.CallbackContext, req *a2a.SendMessageRequest) (*session.Event, error) { t.Fatalf("not called") return nil, nil }, @@ -795,10 +802,10 @@ func TestRemoteAgent_RequestCallbacks(t *testing.T) { { name: "after interceptor short-circuit", after: []AfterA2ARequestCallback{ - func(ctx agent.CallbackContext, req *a2a.MessageSendParams, result *session.Event, err error) (*session.Event, error) { + func(ctx agent.CallbackContext, req *a2a.SendMessageRequest, result *session.Event, err error) (*session.Event, error) { return nil, fmt.Errorf("failed") }, - func(ctx agent.CallbackContext, req *a2a.MessageSendParams, result *session.Event, err error) (*session.Event, error) { + func(ctx agent.CallbackContext, req *a2a.SendMessageRequest, result *session.Event, err error) (*session.Event, error) { t.Fatalf("not called") return nil, nil }, @@ -809,7 +816,7 @@ func TestRemoteAgent_RequestCallbacks(t *testing.T) { name: "after interceptor for empty session", sessionEvents: []*session.Event{}, after: []AfterA2ARequestCallback{ - func(ctx agent.CallbackContext, req *a2a.MessageSendParams, result *session.Event, err error) (*session.Event, error) { + func(ctx agent.CallbackContext, req *a2a.SendMessageRequest, result *session.Event, err error) (*session.Event, error) { if len(req.Message.Parts) != 0 { t.Fatalf("got %d parts, expected empty message", len(req.Message.Parts)) } @@ -820,14 +827,14 @@ func TestRemoteAgent_RequestCallbacks(t *testing.T) { }, { name: "converter error", - converter: func(ctx agent.InvocationContext, req *a2a.MessageSendParams, event a2a.Event, err error) (*session.Event, error) { + converter: func(ctx agent.InvocationContext, req *a2a.SendMessageRequest, event a2a.Event, err error) (*session.Event, error) { return nil, fmt.Errorf("failed") }, wantErr: fmt.Errorf("failed"), }, { name: "converter custom response", - converter: func(ctx agent.InvocationContext, req *a2a.MessageSendParams, event a2a.Event, err error) (*session.Event, error) { + converter: func(ctx agent.InvocationContext, req *a2a.SendMessageRequest, event a2a.Event, err error) (*session.Event, error) { return &session.Event{LLMResponse: model.LLMResponse{Content: genai.NewContentFromText("hello", genai.RoleModel)}}, nil }, wantResponses: []model.LLMResponse{{Content: genai.NewContentFromText("hello", genai.RoleModel)}}, @@ -835,12 +842,12 @@ func TestRemoteAgent_RequestCallbacks(t *testing.T) { { name: "after interceptor invoked with before result", before: []BeforeA2ARequestCallback{ - func(ctx agent.CallbackContext, req *a2a.MessageSendParams) (*session.Event, error) { + func(ctx agent.CallbackContext, req *a2a.SendMessageRequest) (*session.Event, error) { return nil, fmt.Errorf("before error") }, }, after: []AfterA2ARequestCallback{ - func(ctx agent.CallbackContext, req *a2a.MessageSendParams, result *session.Event, err error) (*session.Event, error) { + func(ctx agent.CallbackContext, req *a2a.SendMessageRequest, result *session.Event, err error) (*session.Event, error) { return nil, fmt.Errorf("after error") }, }, @@ -851,20 +858,27 @@ func TestRemoteAgent_RequestCallbacks(t *testing.T) { for _, tc := range testCases { t.Run(tc.name, func(t *testing.T) { executor := &mockA2AExecutor{ - executeFn: func(ctx context.Context, reqCtx *a2asrv.RequestContext, queue eventqueue.Queue) error { - if tc.events != nil { - for _, event := range tc.events(reqCtx) { - if err := queue.Write(ctx, event); err != nil { - return err + executeFn: func(ctx context.Context, execCtx *a2asrv.ExecutorContext) iter.Seq2[a2a.Event, error] { + return func(yield func(a2a.Event, error) bool) { + if tc.events != nil { + for _, event := range tc.events(execCtx) { + if !yield(event, nil) { + return + } } + return } - return nil + yield(a2a.NewMessage(a2a.MessageRoleAgent, a2a.NewTextPart("Hi!")), nil) } - return queue.Write(ctx, a2a.NewMessage(a2a.MessageRoleAgent, a2a.TextPart{Text: "Hi!"})) }, } server := startA2AServer(executor) - card := &a2a.AgentCard{PreferredTransport: a2a.TransportProtocolJSONRPC, URL: server.URL, Capabilities: a2a.AgentCapabilities{Streaming: true}} + card := &a2a.AgentCard{ + SupportedInterfaces: []*a2a.AgentInterface{ + a2a.NewAgentInterface(server.URL, a2a.TransportProtocolJSONRPC), + }, + Capabilities: a2a.AgentCapabilities{Streaming: true}, + } remoteAgent, err := NewA2A(A2AConfig{ Name: "a2a", AgentCard: card, @@ -901,15 +915,15 @@ func TestRemoteAgent_RequestPayload(t *testing.T) { testCases := []struct { name string sessionEvents []*session.Event - wantRequest *a2a.MessageSendParams + wantRequest *a2a.SendMessageRequest }{ { name: "only user message", sessionEvents: []*session.Event{newUserHello()}, - wantRequest: &a2a.MessageSendParams{ + wantRequest: &a2a.SendMessageRequest{ Message: &a2a.Message{ Role: a2a.MessageRoleUser, - Parts: []a2a.Part{a2a.TextPart{Text: "hello"}}, + Parts: a2a.ContentParts{a2a.NewTextPart("hello")}, }, }, }, @@ -930,14 +944,14 @@ func TestRemoteAgent_RequestPayload(t *testing.T) { }, }, }, - wantRequest: &a2a.MessageSendParams{ + wantRequest: &a2a.SendMessageRequest{ Message: &a2a.Message{ Role: a2a.MessageRoleUser, - Parts: []a2a.Part{ - a2a.TextPart{Text: "hello"}, - a2a.TextPart{Text: "For context:"}, - a2a.TextPart{Text: fmt.Sprintf("[%s] said: hi", notRemoteAgentName)}, - a2a.TextPart{Text: "how are you?"}, + Parts: a2a.ContentParts{ + a2a.NewTextPart("hello"), + a2a.NewTextPart("For context:"), + a2a.NewTextPart(fmt.Sprintf("[%s] said: hi", notRemoteAgentName)), + a2a.NewTextPart("how are you?"), }, }, }, @@ -959,14 +973,14 @@ func TestRemoteAgent_RequestPayload(t *testing.T) { {Author: "user", LLMResponse: model.LLMResponse{Content: genai.NewContentFromText("msg3", genai.RoleUser)}}, {Author: notRemoteAgentName, LLMResponse: model.LLMResponse{Content: genai.NewContentFromText("resp3", genai.RoleModel)}}, }, - wantRequest: &a2a.MessageSendParams{ + wantRequest: &a2a.SendMessageRequest{ Message: &a2a.Message{ Role: a2a.MessageRoleUser, ContextID: "ctx-123", - Parts: []a2a.Part{ - a2a.TextPart{Text: "msg3"}, - a2a.TextPart{Text: "For context:"}, - a2a.TextPart{Text: fmt.Sprintf("[%s] said: resp3", notRemoteAgentName)}, + Parts: a2a.ContentParts{ + a2a.NewTextPart("msg3"), + a2a.NewTextPart("For context:"), + a2a.NewTextPart(fmt.Sprintf("[%s] said: resp3", notRemoteAgentName)), }, }, }, @@ -1004,21 +1018,22 @@ func TestRemoteAgent_RequestPayload(t *testing.T) { }, }, }, - wantRequest: &a2a.MessageSendParams{ + wantRequest: &a2a.SendMessageRequest{ Message: &a2a.Message{ Role: a2a.MessageRoleUser, TaskID: "task-1", ContextID: "ctx-1", - Parts: []a2a.Part{ - a2a.TextPart{Text: "lgtm:"}, - a2a.DataPart{ - Data: map[string]any{ + Parts: a2a.ContentParts{ + a2a.NewTextPart("lgtm:"), + func() *a2a.Part { + p := a2a.NewDataPart(map[string]any{ "id": "call-1", "name": "fn", "response": map[string]any{"status": "approved"}, - }, - Metadata: map[string]any{"adk_type": "function_response"}, - }, + }) + p.SetMeta(adka2a.ToA2AMetaKey("type"), "function_response") + return p + }(), }, }, }, @@ -1026,19 +1041,19 @@ func TestRemoteAgent_RequestPayload(t *testing.T) { } server := startA2AServer(newA2AEventReplay(t, []a2a.Event{})) - card := &a2a.AgentCard{PreferredTransport: a2a.TransportProtocolJSONRPC, URL: server.URL, Capabilities: a2a.AgentCapabilities{Streaming: true}} + card := &a2a.AgentCard{SupportedInterfaces: []*a2a.AgentInterface{a2a.NewAgentInterface(server.URL, a2a.TransportProtocolJSONRPC)}, Capabilities: a2a.AgentCapabilities{Streaming: true}} for _, tc := range testCases { t.Run(tc.name, func(t *testing.T) { t.Parallel() errRejected := errors.New("rejected") - var gotRequest *a2a.MessageSendParams + var gotRequest *a2a.SendMessageRequest remoteAgent, err := NewA2A(A2AConfig{ Name: remoteAgentName, AgentCard: card, BeforeRequestCallbacks: []BeforeA2ARequestCallback{ - func(ctx agent.CallbackContext, req *a2a.MessageSendParams) (*session.Event, error) { + func(ctx agent.CallbackContext, req *a2a.SendMessageRequest) (*session.Event, error) { gotRequest = req return nil, errRejected }, @@ -1067,7 +1082,7 @@ func TestRemoteAgent_EmptyResultForEmptySession(t *testing.T) { ictx := newInvocationContext(t, []*session.Event{}) executor := newA2AEventReplay(t, []a2a.Event{ - a2a.NewMessage(a2a.MessageRoleAgent, a2a.TextPart{Text: "will not be invoked, because input is empty"}), + a2a.NewMessage(a2a.MessageRoleAgent, a2a.NewTextPart("will not be invoked, because input is empty")), }) agentName := "a2a agent" @@ -1094,7 +1109,7 @@ func TestRemoteAgent_EmptyResultForEmptySession(t *testing.T) { } func TestRemoteAgent_ResolvesAgentCard(t *testing.T) { - remoteEvents := []a2a.Event{a2a.NewMessage(a2a.MessageRoleAgent, a2a.TextPart{Text: "Hello!"})} + remoteEvents := []a2a.Event{a2a.NewMessage(a2a.MessageRoleAgent, a2a.NewTextPart("Hello!"))} wantResponses := []model.LLMResponse{{Content: genai.NewContentFromText("Hello!", genai.RoleModel), TurnComplete: true}} executor := newA2AEventReplay(t, remoteEvents) @@ -1105,14 +1120,14 @@ func TestRemoteAgent_ResolvesAgentCard(t *testing.T) { mux.Handle("/invoke", a2asrv.NewJSONRPCHandler(handler)) mux.HandleFunc("/.well-known/agent-card.json", func(w http.ResponseWriter, r *http.Request) { url := fmt.Sprintf("%s/invoke", cardServer.URL) - card := &a2a.AgentCard{PreferredTransport: a2a.TransportProtocolJSONRPC, URL: url, Capabilities: a2a.AgentCapabilities{Streaming: true}} + card := &a2a.AgentCard{SupportedInterfaces: []*a2a.AgentInterface{a2a.NewAgentInterface(url, a2a.TransportProtocolJSONRPC)}, Capabilities: a2a.AgentCapabilities{Streaming: true}} if err := json.NewEncoder(w).Encode(card); err != nil { t.Errorf("json.Encode(agentCard) error = %v", err) } }) cardServer = httptest.NewServer(mux) - remoteAgent, err := NewA2A(A2AConfig{Name: "a2a", AgentCardSource: cardServer.URL}) + remoteAgent, err := NewA2A(A2AConfig{Name: "a2a", AgentCardProvider: NewAgentCardProvider(cardServer.URL)}) if err != nil { t.Fatalf("remoteagent.NewA2A() error = %v", err) } @@ -1133,17 +1148,23 @@ func TestRemoteAgent_ResolvesAgentCard(t *testing.T) { } func TestRemoteAgent_ErrorEventIfNoCompatibleTransport(t *testing.T) { - remoteEvents := []a2a.Event{a2a.NewMessage(a2a.MessageRoleAgent, a2a.TextPart{Text: "will not be invoked!"})} + remoteEvents := []a2a.Event{a2a.NewMessage(a2a.MessageRoleAgent, a2a.NewTextPart("will not be invoked!"))} executor := newA2AEventReplay(t, remoteEvents) server := startA2AServer(executor) - remoteAgent, err := NewA2A(A2AConfig{ - Name: "a2a", - ClientFactory: a2aclient.NewFactory(a2aclient.WithDefaultsDisabled()), - AgentCard: &a2a.AgentCard{ - PreferredTransport: a2a.TransportProtocolJSONRPC, - URL: server.URL, + name := "a2a" + cardResource := &a2a.AgentCard{ + SupportedInterfaces: []*a2a.AgentInterface{ + a2a.NewAgentInterface(server.URL, a2a.TransportProtocolJSONRPC), }, + Name: name, + } + remoteAgent, err := NewA2A(A2AConfig{ + Name: name, + AgentCard: cardResource, + ClientProvider: NewA2AClientProvider( + a2aclient.NewFactory(a2aclient.WithDefaultsDisabled()), + ), }) if err != nil { t.Fatalf("remoteagent.NewA2A() error = %v", err) @@ -1166,11 +1187,12 @@ func TestRemoteAgent_ErrorEventIfNoCompatibleTransport(t *testing.T) { func TestRemoteAgent_ErrorEventOnServerError(t *testing.T) { executorErr := fmt.Errorf("mockExecutor failed") executor := &mockA2AExecutor{ - executeFn: func(ctx context.Context, reqCtx *a2asrv.RequestContext, q eventqueue.Queue) error { - return executorErr + executeFn: func(ctx context.Context, execCtx *a2asrv.ExecutorContext) iter.Seq2[a2a.Event, error] { + return func(yield func(a2a.Event, error) bool) { + yield(nil, executorErr) + } }, } - remoteAgent := newA2ARemoteAgent(t, "a2a agent", startA2AServer(executor)) ictx := newInvocationContext(t, []*session.Event{newUserHello()}) @@ -1188,16 +1210,16 @@ func TestRemoteAgent_ErrorEventOnServerError(t *testing.T) { } func TestRemoteAgent_CustomConverters(t *testing.T) { - originalA2APart := a2a.TextPart{Text: "hello"} - customA2APart := a2a.TextPart{Text: "modified"} - mockGenAIPartConverter := func(ctx context.Context, event *session.Event, part *genai.Part) (a2a.Part, error) { + originalA2APart := a2a.NewTextPart("hello") + customA2APart := a2a.NewTextPart("modified") + mockGenAIPartConverter := func(ctx context.Context, event *session.Event, part *genai.Part) (*a2a.Part, error) { return customA2APart, nil } tests := []struct { name string cfg A2AConfig - want a2a.Part + want *a2a.Part }{ { name: "custom converter", @@ -1219,7 +1241,7 @@ func TestRemoteAgent_CustomConverters(t *testing.T) { if len(msg.Parts) != 1 { t.Fatalf("len(msg.Parts) = %d, want 1", len(msg.Parts)) } - if textPart, ok := msg.Parts[0].(a2a.TextPart); !ok || textPart.Text != tc.want.(a2a.TextPart).Text { + if msg.Parts[0].Text() != tc.want.Text() { t.Fatalf("msg.Parts[0] = %+v, want %+v", msg.Parts[0], tc.want) } } @@ -1228,7 +1250,7 @@ func TestRemoteAgent_CustomConverters(t *testing.T) { func TestRemoteAgent_CleanupCallback(t *testing.T) { testCases := []struct { name string - events func(*a2asrv.RequestContext) []a2a.Event + events func(*a2asrv.ExecutorContext) []a2a.Event afterRequestCallbacks []AfterA2ARequestCallback eventConverter A2AEventConverter breakAfter int @@ -1238,7 +1260,7 @@ func TestRemoteAgent_CleanupCallback(t *testing.T) { { name: "after request callback error", afterRequestCallbacks: []AfterA2ARequestCallback{ - func(ctx agent.CallbackContext, req *a2a.MessageSendParams, resp *session.Event, err error) (*session.Event, error) { + func(ctx agent.CallbackContext, req *a2a.SendMessageRequest, resp *session.Event, err error) (*session.Event, error) { return nil, fmt.Errorf("callback error") }, }, @@ -1246,7 +1268,7 @@ func TestRemoteAgent_CleanupCallback(t *testing.T) { }, { name: "part converter error", - eventConverter: func(ctx agent.InvocationContext, req *a2a.MessageSendParams, event a2a.Event, err error) (*session.Event, error) { + eventConverter: func(ctx agent.InvocationContext, req *a2a.SendMessageRequest, event a2a.Event, err error) (*session.Event, error) { if _, ok := event.(*a2a.TaskArtifactUpdateEvent); ok { return nil, fmt.Errorf("converter error") } @@ -1273,38 +1295,46 @@ func TestRemoteAgent_CleanupCallback(t *testing.T) { cleanupTaskID a2a.TaskID cleanupCause error ) - cleanupCallback := func(ctx context.Context, card *a2a.AgentCard, client *a2aclient.Client, task a2a.TaskInfo, cause error) { + cleanupCallback := func(ctx context.Context, card *a2a.AgentCard, client A2AClient, task a2a.TaskInfo, cause error) { cleanupCalled = true cleanupTaskID = task.TaskID cleanupCause = cause - if _, err := client.CancelTask(ctx, &a2a.TaskIDParams{ID: task.TaskID}); err != nil { + if _, err := client.CancelTask(ctx, &a2a.CancelTaskRequest{ID: task.TaskID}); err != nil { t.Errorf("client.CancelTask() error = %v", err) } } remoteTaskIDChan := make(chan a2a.TaskID, 1) executor := &mockA2AExecutor{ - executeFn: func(ctx context.Context, reqCtx *a2asrv.RequestContext, queue eventqueue.Queue) error { - remoteTaskIDChan <- reqCtx.TaskID - if err := queue.Write(ctx, a2a.NewSubmittedTask(reqCtx, reqCtx.Message)); err != nil { - return err + cancelFn: func(ctx context.Context, reqCtx *a2asrv.ExecutorContext) iter.Seq2[a2a.Event, error] { + return func(yield func(a2a.Event, error) bool) { + yield(a2a.NewStatusUpdateEvent(reqCtx, a2a.TaskStateCanceled, nil), nil) } - for ctx.Err() == nil { - data := a2a.DataPart{Data: map[string]any{"foo": "bar"}} - if err := queue.Write(ctx, a2a.NewArtifactEvent(reqCtx, data)); err != nil { - return err + }, + executeFn: func(ctx context.Context, reqCtx *a2asrv.ExecutorContext) iter.Seq2[a2a.Event, error] { + return func(yield func(a2a.Event, error) bool) { + remoteTaskIDChan <- reqCtx.TaskID + if !yield(a2a.NewSubmittedTask(reqCtx, reqCtx.Message), nil) { + return + } + for ctx.Err() == nil { + data := a2a.NewDataPart(map[string]any{"foo": "bar"}) + if !yield(a2a.NewArtifactEvent(reqCtx, data), nil) { + return + } + time.Sleep(1 * time.Millisecond) } - time.Sleep(1 * time.Millisecond) + yield(a2a.NewStatusUpdateEvent(reqCtx, a2a.TaskStateCompleted, nil), nil) } - finalUpdate := a2a.NewStatusUpdateEvent(reqCtx, a2a.TaskStateCompleted, nil) - finalUpdate.Final = true - return queue.Write(ctx, finalUpdate) }, } server := startA2AServer(executor) defer server.Close() - card := &a2a.AgentCard{PreferredTransport: a2a.TransportProtocolJSONRPC, URL: server.URL, Capabilities: a2a.AgentCapabilities{Streaming: true}} + card := &a2a.AgentCard{ + SupportedInterfaces: []*a2a.AgentInterface{a2a.NewAgentInterface(server.URL, a2a.TransportProtocolJSONRPC)}, + Capabilities: a2a.AgentCapabilities{Streaming: true}, + } remoteAgent, err := NewA2A(A2AConfig{ Name: "a2a", AgentCard: card, @@ -1356,7 +1386,7 @@ func TestRemoteAgent_CleanupCallback(t *testing.T) { } client := newA2AClient(t, server) - task, err := client.GetTask(t.Context(), &a2a.TaskQueryParams{ID: expectedTaskID}) + task, err := client.GetTask(t.Context(), &a2a.GetTaskRequest{ID: expectedTaskID}) if err != nil { t.Fatalf("client.CancelTask() error = %v", err) } @@ -1376,11 +1406,11 @@ func TestRemoteAgent_PartConverter(t *testing.T) { } cfg := A2AConfig{ - GenAIPartConverter: func(ctx context.Context, event *session.Event, p *genai.Part) (a2a.Part, error) { + GenAIPartConverter: func(ctx context.Context, event *session.Event, p *genai.Part) (*a2a.Part, error) { if p.Text == "DROP" { return nil, nil } - return a2a.TextPart{Text: p.Text}, nil + return a2a.NewTextPart(p.Text), nil }, } @@ -1400,8 +1430,8 @@ func TestRemoteAgent_PartConverter(t *testing.T) { t.Fatalf("got nil part, want it filtered out.") } - if tp, ok := p.(a2a.TextPart); ok && tp.Text != "KEEP" { - t.Errorf("got %s, want 'KEEP'", tp.Text) + if p.Text() != "KEEP" { + t.Errorf("got %s, want 'KEEP'", p.Text()) } } } diff --git a/agent/remoteagent/a2a_e2e_test.go b/agent/remoteagent/v2/a2a_e2e_test.go similarity index 81% rename from agent/remoteagent/a2a_e2e_test.go rename to agent/remoteagent/v2/a2a_e2e_test.go index 4799bd60d..7421a23e5 100644 --- a/agent/remoteagent/a2a_e2e_test.go +++ b/agent/remoteagent/v2/a2a_e2e_test.go @@ -29,10 +29,9 @@ import ( "testing" "time" - "github.com/a2aproject/a2a-go/a2a" - "github.com/a2aproject/a2a-go/a2aclient" - "github.com/a2aproject/a2a-go/a2asrv" - "github.com/a2aproject/a2a-go/a2asrv/eventqueue" + "github.com/a2aproject/a2a-go/v2/a2a" + "github.com/a2aproject/a2a-go/v2/a2aclient" + "github.com/a2aproject/a2a-go/v2/a2asrv" "github.com/google/go-cmp/cmp" "github.com/google/go-cmp/cmp/cmpopts" "google.golang.org/genai" @@ -47,7 +46,7 @@ import ( "google.golang.org/adk/model" "google.golang.org/adk/model/gemini" "google.golang.org/adk/runner" - "google.golang.org/adk/server/adka2a" + "google.golang.org/adk/server/adka2a/v2" "google.golang.org/adk/session" "google.golang.org/adk/tool" "google.golang.org/adk/tool/functiontool" @@ -94,10 +93,10 @@ func TestA2AInputRequired(t *testing.T) { return createLongRunningToolApproval(t, pendingResponse) }, wantFirstArtifactParts: a2a.ContentParts{ - a2a.TextPart{Text: modelTextRequiresApproval}, - a2a.TextPart{Text: modelTextWaitingForApproval}, + a2a.NewTextPart(modelTextRequiresApproval), + a2a.NewTextPart(modelTextWaitingForApproval), }, - wantSecondArtifactParts: a2a.ContentParts{a2a.TextPart{Text: modelTextTaskComplete}}, + wantSecondArtifactParts: a2a.ContentParts{a2a.NewTextPart(modelTextTaskComplete)}, }, { name: "tool confirmation", @@ -106,28 +105,32 @@ func TestA2AInputRequired(t *testing.T) { return createToolConfirmationApproval(t, toolCall) }, wantFirstArtifactParts: a2a.ContentParts{ - a2a.TextPart{Text: modelTextRequiresApproval}, - a2a.DataPart{ - Data: map[string]any{"name": approvalToolName}, - Metadata: map[string]any{"adk_is_long_running": false, "adk_type": "function_call"}, - }, - a2a.DataPart{ - Data: map[string]any{ + a2a.NewTextPart(modelTextRequiresApproval), + func() *a2a.Part { + p := a2a.NewDataPart(map[string]any{"name": approvalToolName}) + p.SetMeta(adka2a.ToA2AMetaKey("is_long_running"), false) + p.SetMeta(adka2a.ToA2AMetaKey("type"), "function_call") + return p + }(), + func() *a2a.Part { + p := a2a.NewDataPart(map[string]any{ "name": approvalToolName, "response": map[string]any{"status": string(approvalStatusPending)}, - }, - Metadata: map[string]any{"adk_type": "function_response"}, - }, + }) + p.SetMeta(adka2a.ToA2AMetaKey("type"), "function_response") + return p + }(), }, wantSecondArtifactParts: a2a.ContentParts{ - a2a.DataPart{ - Data: map[string]any{ + func() *a2a.Part { + p := a2a.NewDataPart(map[string]any{ "name": approvalToolName, "response": map[string]any{"status": string(approvalStatusVerified)}, - }, - Metadata: map[string]any{"adk_type": "function_response"}, - }, - a2a.TextPart{Text: modelTextTaskComplete}, + }) + p.SetMeta(adka2a.ToA2AMetaKey("type"), "function_response") + return p + }(), + a2a.NewTextPart(modelTextTaskComplete), }, }, } @@ -147,7 +150,7 @@ func TestA2AInputRequired(t *testing.T) { // Initial message triggers input required taskContent := "Perform important task!" - msg1 := a2a.NewMessage(a2a.MessageRoleUser, a2a.TextPart{Text: taskContent}) + msg1 := a2a.NewMessage(a2a.MessageRoleUser, a2a.NewTextPart(taskContent)) task1 := mustSendMessage(t, client, msg1) if task1.Status.State != a2a.TaskStateInputRequired { t.Fatalf("client.SendMessage(Initial) result state = %q, want %q", task1.Status.State, a2a.TaskStateInputRequired) @@ -158,7 +161,7 @@ func TestA2AInputRequired(t *testing.T) { // Incomplete followup keeps the task in input-required incompleteFollowupText := "Is it really necessary?" - msg2 := a2a.NewMessageForTask(a2a.MessageRoleUser, task1, a2a.TextPart{Text: incompleteFollowupText}) + msg2 := a2a.NewMessageForTask(a2a.MessageRoleUser, task1, a2a.NewTextPart(incompleteFollowupText)) task2 := mustSendMessage(t, client, msg2) if task2.Status.State != a2a.TaskStateInputRequired { t.Fatalf("client.SendMessage(IncompleteInput) result state = %q, want %q", task2.Status.State, a2a.TaskStateInputRequired) @@ -175,16 +178,16 @@ func TestA2AInputRequired(t *testing.T) { } // The last part should be the error message lastPart := task2.Status.Message.Parts[len(task2.Status.Message.Parts)-1] - tp, ok := lastPart.(a2a.TextPart) - if !ok { + text := lastPart.Text() + if text == "" { t.Fatalf("last part is not TextPart") } - if !strings.Contains(tp.Text, "no input provided") { - t.Errorf("last part text = %q; want it to contain 'no input provided'", tp.Text) + if !strings.Contains(text, "no input provided") { + t.Errorf("last part text = %q; want it to contain 'no input provided'", text) } // Another incomplete followup should not accumulate error messages - msg2a := a2a.NewMessageForTask(a2a.MessageRoleUser, task1, a2a.TextPart{Text: "Still debating?"}) + msg2a := a2a.NewMessageForTask(a2a.MessageRoleUser, task1, a2a.NewTextPart("Still debating?")) task2a := mustSendMessage(t, client, msg2a) if task2a.Status.State != a2a.TaskStateInputRequired { t.Fatalf("client.SendMessage(IncompleteInput 2) result state = %q, want %q", task2a.Status.State, a2a.TaskStateInputRequired) @@ -193,7 +196,7 @@ func TestA2AInputRequired(t *testing.T) { // Count validation errors in parts validationErrors := 0 for _, p := range task2a.Status.Message.Parts { - if tp, ok := p.(a2a.TextPart); ok && strings.Contains(tp.Text, "no input provided") { + if strings.Contains(p.Text(), "no input provided") { validationErrors++ } } @@ -206,7 +209,7 @@ func TestA2AInputRequired(t *testing.T) { approvedResponse := tc.createApproval(t, toolCall, pendingResponse) msg3 := a2a.NewMessageForTask(a2a.MessageRoleUser, task2, - a2a.TextPart{Text: "LGTM"}, + a2a.NewTextPart("LGTM"), toA2AParts(t, []*genai.Part{approvedResponse}, []string{toolCall.ID})[0], ) task3 := mustSendMessage(t, client, msg3) @@ -270,7 +273,7 @@ func TestA2AMultiHopInputRequired(t *testing.T) { genai.NewPartFromText(modelTextWaitingForApproval), }, []string{}), wantSecondArtifactParts: a2a.ContentParts{ - a2a.TextPart{Text: modelTextTaskComplete}, + a2a.NewTextPart(modelTextTaskComplete), }, }, { @@ -288,14 +291,15 @@ func TestA2AMultiHopInputRequired(t *testing.T) { genai.NewPartFromFunctionResponse(approvalToolName, map[string]any{"status": string(approvalStatusPending)}), }, []string{}), wantSecondArtifactParts: a2a.ContentParts{ - a2a.DataPart{ - Data: map[string]any{ + func() *a2a.Part { + p := a2a.NewDataPart(map[string]any{ "name": approvalToolName, "response": map[string]any{"status": string(approvalStatusVerified)}, - }, - Metadata: map[string]any{"adk_type": "function_response"}, - }, - a2a.TextPart{Text: modelTextTaskComplete}, + }) + p.SetMeta(adka2a.ToA2AMetaKey("type"), "function_response") + return p + }(), + a2a.NewTextPart(modelTextTaskComplete), }, }, } @@ -321,14 +325,14 @@ func TestA2AMultiHopInputRequired(t *testing.T) { client := newA2AClient(t, serverA) // Initial message triggers input required - msg1 := a2a.NewMessage(a2a.MessageRoleUser, a2a.TextPart{Text: "Hello, perform important task!"}) + msg1 := a2a.NewMessage(a2a.MessageRoleUser, a2a.NewTextPart("Hello, perform important task!")) task1 := mustSendMessage(t, client, msg1) if task1.Status.State != a2a.TaskStateInputRequired { t.Fatalf("client.SendMessage(Initial) result state = %q, want %q", task1.Status.State, a2a.TaskStateInputRequired) } // Incomplete followup keeps the task in input-required - msg2 := a2a.NewMessageForTask(a2a.MessageRoleUser, task1, a2a.TextPart{Text: "Is it really necessary?"}) + msg2 := a2a.NewMessageForTask(a2a.MessageRoleUser, task1, a2a.NewTextPart("Is it really necessary?")) task2 := mustSendMessage(t, client, msg2) if task2.Status.State != a2a.TaskStateInputRequired { t.Fatalf("client.SendMessage(IncompleteInput) result state = %q, want %q", task2.Status.State, a2a.TaskStateInputRequired) @@ -338,7 +342,7 @@ func TestA2AMultiHopInputRequired(t *testing.T) { toolCall, pendingResponse := findLongRunningCall(t, toGenaiParts(t, filterPartial(task2.Status.Message.Parts))) approvedResponse := tc.createApproval(t, toolCall, pendingResponse) msg3 := a2a.NewMessageForTask(a2a.MessageRoleUser, task2, - a2a.TextPart{Text: "LGTM"}, + a2a.NewTextPart("LGTM"), toA2AParts(t, []*genai.Part{approvedResponse}, nil)[0], ) task3 := mustSendMessage(t, client, msg3) @@ -377,22 +381,27 @@ func TestA2ACleanupPropagation(t *testing.T) { // until it detects a context cancelation remoteTaskIDChan, remoteCleanupCalledChan := make(chan a2a.TaskID, 1), make(chan struct{}, 2) serverB := startA2AServer(&mockA2AExecutor{ - executeFn: func(ctx context.Context, reqCtx *a2asrv.RequestContext, queue eventqueue.Queue) error { - remoteTaskIDChan <- reqCtx.TaskID - if err := queue.Write(ctx, a2a.NewSubmittedTask(reqCtx, reqCtx.Message)); err != nil { - return err + cancelFn: func(ctx context.Context, reqCtx *a2asrv.ExecutorContext) iter.Seq2[a2a.Event, error] { + return func(yield func(a2a.Event, error) bool) { + yield(a2a.NewStatusUpdateEvent(reqCtx, a2a.TaskStateCanceled, nil), nil) } - for ctx.Err() == nil { - if err := queue.Write(ctx, a2a.NewArtifactEvent(reqCtx, a2a.TextPart{Text: "foo"})); err != nil { - return err + }, + executeFn: func(ctx context.Context, reqCtx *a2asrv.ExecutorContext) iter.Seq2[a2a.Event, error] { + return func(yield func(a2a.Event, error) bool) { + remoteTaskIDChan <- reqCtx.TaskID + if !yield(a2a.NewSubmittedTask(reqCtx, reqCtx.Message), nil) { + return + } + for ctx.Err() == nil { + if !yield(a2a.NewArtifactEvent(reqCtx, a2a.NewTextPart("foo")), nil) { + return + } + time.Sleep(1 * time.Millisecond) } - time.Sleep(1 * time.Millisecond) + yield(a2a.NewStatusUpdateEvent(reqCtx, a2a.TaskStateCompleted, nil), nil) } - finalUpdate := a2a.NewStatusUpdateEvent(reqCtx, a2a.TaskStateCompleted, nil) - finalUpdate.Final = true - return queue.Write(ctx, finalUpdate) }, - cleanupFn: func(ctx context.Context, reqCtx *a2asrv.RequestContext, result a2a.SendMessageResult, cause error) { + cleanupFn: func(ctx context.Context, reqCtx *a2asrv.ExecutorContext, result a2a.SendMessageResult, cause error) { remoteCleanupCalledChan <- struct{}{} }, }) @@ -401,7 +410,18 @@ func TestA2ACleanupPropagation(t *testing.T) { // Root server connects to server B through remote subagent remoteAgentB := newA2ARemoteAgent(t, "remote-agent-b", serverB) rootA := newRootAgent("agent-b", remoteAgentB) - executorA := newAgentExecutor(rootA, nil, adka2a.OutputArtifactPerEvent) + executorCleanupCalledChan := make(chan struct{}, 2) + executorA := adka2a.NewExecutor(adka2a.ExecutorConfig{ + OutputMode: adka2a.OutputArtifactPerRun, + RunnerConfig: runner.Config{ + AppName: rootA.Name(), + SessionService: session.InMemoryService(), + Agent: rootA, + }, + A2AExecutionCleanupCallback: func(ctx context.Context, reqCtx *a2asrv.ExecutorContext, subAgentCards []*a2a.AgentCard, result a2a.SendMessageResult, cause error) { + executorCleanupCalledChan <- struct{}{} + }, + }) serverA := startA2AServer(executorA) defer serverA.Close() @@ -411,8 +431,8 @@ func TestA2ACleanupPropagation(t *testing.T) { statusUpdateEventChan := make(chan a2a.Event, 10) go func() { defer close(statusUpdateEventChan) - msg := a2a.NewMessage(a2a.MessageRoleUser, a2a.TextPart{Text: "work"}) - for event, err := range client.SendStreamingMessage(t.Context(), &a2a.MessageSendParams{Message: msg}) { + msg := a2a.NewMessage(a2a.MessageRoleUser, a2a.NewTextPart("work")) + for event, err := range client.SendStreamingMessage(t.Context(), &a2a.SendMessageRequest{Message: msg}) { if err != nil { t.Errorf("client.SendStreamingMessage() error = %v", err) return @@ -429,7 +449,7 @@ func TestA2ACleanupPropagation(t *testing.T) { cancelResultChan := make(chan *a2a.Task, 1) go func() { defer close(cancelResultChan) - task, err := client.CancelTask(t.Context(), &a2a.TaskIDParams{ID: taskID}) + task, err := client.CancelTask(t.Context(), &a2a.CancelTaskRequest{ID: taskID}) if err != nil { t.Errorf("client.CancelTask() error = %v", err) return @@ -452,11 +472,25 @@ func TestA2ACleanupPropagation(t *testing.T) { // Check subagent task got cancelled when the parent task was cancelled. // Reads from channel twice because cleanup gets called both for cancelation and execution. - <-remoteCleanupCalledChan - <-remoteCleanupCalledChan + timeout := time.After(5 * time.Second) + for range 2 { + select { + case <-remoteCleanupCalledChan: + case <-timeout: + t.Fatalf("remote cleanup was not called") + } + } remoteTaskID := <-remoteTaskIDChan + + for range 2 { + select { + case <-executorCleanupCalledChan: + case <-timeout: + t.Fatalf("executor cleanup was not called") + } + } remoteClient := newA2AClient(t, serverB) - remoteTask, err := remoteClient.GetTask(t.Context(), &a2a.TaskQueryParams{ID: remoteTaskID}) + remoteTask, err := remoteClient.GetTask(t.Context(), &a2a.GetTaskRequest{ID: remoteTaskID}) if err != nil { t.Fatalf("remoteClient.GetTask() error = %v", err) } @@ -494,8 +528,8 @@ func TestA2ASingleHopFinalResponse(t *testing.T) { }, wantState: a2a.TaskStateCompleted, wantArtifactParts: a2a.ContentParts{ - a2a.TextPart{Text: "Hello, I am beep!"}, - a2a.TextPart{Text: "I am boop. We are here to help!"}, + a2a.NewTextPart("Hello, I am beep!"), + a2a.NewTextPart("I am boop. We are here to help!"), }, wantPartial: true, }, @@ -514,8 +548,8 @@ func TestA2ASingleHopFinalResponse(t *testing.T) { }, wantState: a2a.TaskStateCompleted, wantArtifactParts: a2a.ContentParts{ - a2a.TextPart{Text: "Hello, I am beep!"}, - a2a.TextPart{Text: "I am boop. We are here to help!"}, + a2a.NewTextPart("Hello, I am beep!"), + a2a.NewTextPart("I am boop. We are here to help!"), }, }, { @@ -589,7 +623,7 @@ func TestA2ASingleHopFinalResponse(t *testing.T) { defer server.Close() client := newA2AClient(t, server) - msg := a2a.NewMessage(a2a.MessageRoleUser, a2a.TextPart{Text: "Tell me about the current weather"}) + msg := a2a.NewMessage(a2a.MessageRoleUser, a2a.NewTextPart("Tell me about the current weather")) task := mustSendMessage(t, client, msg) if task.Status.State != tc.wantState { t.Fatalf("client.SendMessage(Initial) result state = %q, want %q", task.Status.State, tc.wantState) @@ -610,7 +644,7 @@ func TestA2ASingleHopFinalResponse(t *testing.T) { if task.Status.Message == nil || len(task.Status.Message.Parts) != 1 { t.Fatalf("got status message = %v, want message with one part", task.Status.Message) } - if tp, ok := task.Status.Message.Parts[0].(a2a.TextPart); !ok || !strings.Contains(tp.Text, tc.wantStatusContain) { + if !strings.Contains(task.Status.Message.Parts[0].Text(), tc.wantStatusContain) { t.Fatalf("got status message = %v, want text containing %q", task.Status.Message.Parts[0], tc.wantStatusContain) } } @@ -635,7 +669,13 @@ func TestA2ASingleHopFinalResponse(t *testing.T) { } else { partialArtifact = task.Artifacts[1] } - wantPartialParts := a2a.ContentParts{a2a.DataPart{Data: map[string]any{}, Metadata: map[string]any{"adk_partial": true}}} + wantPartialParts := a2a.ContentParts{ + func() *a2a.Part { + p := a2a.NewDataPart(map[string]any{}) + p.SetMeta(adka2a.ToA2AMetaKey("partial"), true) + return p + }(), + } if diff := cmp.Diff(wantPartialParts, partialArtifact.Parts); diff != "" { t.Fatalf("task wrong artifact parts (+got,-want) diff = %s", diff) } @@ -664,13 +704,13 @@ func TestA2ARemoteAgentStreamingGeminiSuccess(t *testing.T) { ctx := t.Context() client := newA2AClient(t, serverA) - msg := a2a.NewMessage(a2a.MessageRoleUser, a2a.TextPart{Text: "tell me about the capital of Poland"}) + msg := a2a.NewMessage(a2a.MessageRoleUser, a2a.NewTextPart("tell me about the capital of Poland")) msg.ContextID = a2a.NewContextID() // Make streaming request and aggregate results var taskID a2a.TaskID partialText, finalText := "", "" - for event, err := range client.SendStreamingMessage(t.Context(), &a2a.MessageSendParams{Message: msg}) { + for event, err := range client.SendStreamingMessage(t.Context(), &a2a.SendMessageRequest{Message: msg}) { if err != nil { t.Fatalf("client.SendStreamingMessage() error = %v", err) } @@ -679,7 +719,7 @@ func TestA2ARemoteAgentStreamingGeminiSuccess(t *testing.T) { if len(tau.Artifact.Parts) != 1 { t.Fatalf("got %d parts in final partial artifact update, want 1", len(tau.Artifact.Parts)) } - if dp, ok := tau.Artifact.Parts[0].(a2a.DataPart); !ok || len(dp.Data) > 0 { + if dp := tau.Artifact.Parts[0]; dp.Data() == nil || len(dp.Data().(map[string]any)) > 0 { t.Fatalf("got %v part in final partial artifact update, want empty data part", tau.Artifact.Parts[0]) } continue @@ -687,7 +727,7 @@ func TestA2ARemoteAgentStreamingGeminiSuccess(t *testing.T) { if adka2a.IsPartial(tau.Metadata) { for _, p := range tau.Artifact.Parts { - partialText += p.(a2a.TextPart).Text + partialText += p.Text() } continue } @@ -695,7 +735,7 @@ func TestA2ARemoteAgentStreamingGeminiSuccess(t *testing.T) { if len(finalText) > 0 { t.Fatal("got multiple non-partial updates, want 1") } - finalText = tau.Artifact.Parts[0].(a2a.TextPart).Text + finalText = tau.Artifact.Parts[0].Text() } taskID = event.TaskInfo().TaskID } @@ -709,7 +749,7 @@ func TestA2ARemoteAgentStreamingGeminiSuccess(t *testing.T) { } // Check A2A Task state - task, err := client.GetTask(ctx, &a2a.TaskQueryParams{ID: taskID}) + task, err := client.GetTask(ctx, &a2a.GetTaskRequest{ID: taskID}) if err != nil { t.Fatalf("client.GetTask() error = %v", err) } @@ -776,12 +816,12 @@ func TestA2ARemoteAgentStreamingGeminiError(t *testing.T) { ctx := t.Context() client := newA2AClient(t, serverA) - msg := a2a.NewMessage(a2a.MessageRoleUser, a2a.TextPart{Text: "tell me about the capital of Poland"}) + msg := a2a.NewMessage(a2a.MessageRoleUser, a2a.NewTextPart("tell me about the capital of Poland")) msg.ContextID = a2a.NewContextID() // Make streaming request and aggregate results var taskID a2a.TaskID - for event, err := range client.SendStreamingMessage(t.Context(), &a2a.MessageSendParams{Message: msg}) { + for event, err := range client.SendStreamingMessage(t.Context(), &a2a.SendMessageRequest{Message: msg}) { if err != nil { t.Fatalf("client.SendStreamingMessage() error = %v", err) } @@ -789,7 +829,7 @@ func TestA2ARemoteAgentStreamingGeminiError(t *testing.T) { } // Check A2A Task state - task, err := client.GetTask(ctx, &a2a.TaskQueryParams{ID: taskID}) + task, err := client.GetTask(ctx, &a2a.GetTaskRequest{ID: taskID}) if err != nil { t.Fatalf("client.GetTask() error = %v", err) } @@ -799,13 +839,13 @@ func TestA2ARemoteAgentStreamingGeminiError(t *testing.T) { if task.Status.Message == nil || len(task.Status.Message.Parts) != 1 { t.Fatalf("task status message = %v, want 1 part", task.Status.Message) } - if tp, ok := task.Status.Message.Parts[0].(a2a.TextPart); !ok || !strings.Contains(tp.Text, errorMessage) { + if !strings.Contains(task.Status.Message.Parts[0].Text(), errorMessage) { t.Fatalf("task status message = %v, want containing %q", task.Status.Message.Parts[0], errorMessage) } if len(task.Artifacts) != 1 || len(adka2a.WithoutPartialArtifacts(task.Artifacts)) != 0 { t.Fatalf("task artifacts = %v, want single partial artifact", task.Artifacts) } - if dp, ok := task.Artifacts[0].Parts[0].(a2a.DataPart); !ok || len(dp.Data) != 0 { + if dp := task.Artifacts[0].Parts[0]; dp.Data() == nil { t.Fatalf("task artifact = %v, want reset partial artifact", task.Artifacts[0]) } @@ -846,7 +886,7 @@ func newLongRunningTool(t *testing.T) tool.Tool { Name: approvalToolName, Description: "Request approval before proceeding.", IsLongRunning: true, - }, func(ctx tool.Context, x map[string]any) (approval, error) { + }, func(ctx agent.ToolContext, x map[string]any) (approval, error) { return approval{Status: approvalStatusPending, TicketID: a2a.NewContextID()}, nil }) if err != nil { @@ -861,7 +901,7 @@ func newToolConfirmation(t *testing.T) tool.Tool { requestApproval, err := functiontool.New(functiontool.Config{ Name: approvalToolName, Description: "Request approval before proceeding.", - }, func(ctx tool.Context, x map[string]any) (approval, error) { + }, func(ctx agent.ToolContext, x map[string]any) (approval, error) { confirmation := ctx.ToolConfirmation() if confirmation == nil { ticketID := a2a.NewContextID() @@ -955,7 +995,7 @@ func newAgentExecutor(agnt agent.Agent, service session.Service, mode adka2a.Out func mustSendMessage(t *testing.T, client *a2aclient.Client, msg *a2a.Message) *a2a.Task { t.Helper() - sendParams := &a2a.MessageSendParams{Message: msg} + sendParams := &a2a.SendMessageRequest{Message: msg} result, err := client.SendMessage(t.Context(), sendParams) if err != nil { t.Fatalf("client.SendMessage() error = %v", err) @@ -967,8 +1007,8 @@ func mustSendMessage(t *testing.T, client *a2aclient.Client, msg *a2a.Message) * return task } -func filterPartial(parts []a2a.Part) []a2a.Part { - var result []a2a.Part +func filterPartial(parts []*a2a.Part) []*a2a.Part { + var result []*a2a.Part for _, p := range parts { if b, _ := p.Meta()[adka2a.ToA2AMetaKey("partial")].(bool); b { continue @@ -1000,16 +1040,16 @@ func findLongRunningCall(t *testing.T, parts []*genai.Part) (*genai.FunctionCall return call, response } -func toA2AParts(t *testing.T, parts []*genai.Part, callIDs []string) []a2a.Part { +func toA2AParts(t *testing.T, parts []*genai.Part, callIDs []string) []*a2a.Part { t.Helper() - a2aParts, err := adka2a.ToA2AParts(parts, callIDs) + result, err := adka2a.ToA2AParts(parts, callIDs) if err != nil { t.Fatalf("adka2a.ToA2AParts() error = %v", err) } - return a2aParts + return result } -func toGenaiParts(t *testing.T, a2aParts []a2a.Part) []*genai.Part { +func toGenaiParts(t *testing.T, a2aParts []*a2a.Part) []*genai.Part { t.Helper() parts, err := adka2a.ToGenAIParts(a2aParts) if err != nil { @@ -1040,8 +1080,10 @@ func newA2AClient(t *testing.T, server *testA2AServer) *a2aclient.Client { t.Helper() result, err := a2aclient.NewFromCard(t.Context(), &a2a.AgentCard{ - PreferredTransport: a2a.TransportProtocolJSONRPC, - URL: server.URL, Capabilities: a2a.AgentCapabilities{Streaming: true}, + SupportedInterfaces: []*a2a.AgentInterface{ + a2a.NewAgentInterface(server.URL, a2a.TransportProtocolJSONRPC), + }, + Capabilities: a2a.AgentCapabilities{Streaming: true}, }) if err != nil { t.Fatalf("a2aclient.NewFromEndpoints() error = %v", err) @@ -1129,19 +1171,23 @@ func TestA2AMultiHopInputRequiredCancellation(t *testing.T) { remoteAgentName := "remote-agent-B" remoteTaskIDChan := make(chan a2a.TaskID, 1) serverB := startA2AServer(&mockA2AExecutor{ - executeFn: func(ctx context.Context, reqCtx *a2asrv.RequestContext, queue eventqueue.Queue) error { - remoteTaskIDChan <- reqCtx.TaskID - ev := a2a.NewStatusUpdateEvent(reqCtx, a2a.TaskStateInputRequired, a2a.NewMessage(a2a.MessageRoleAgent, a2a.DataPart{ - Data: map[string]any{"id": "call-1", "name": "foo"}, - Metadata: map[string]any{"adk_is_long_running": true, "adk_type": "function_call"}, - })) - ev.Final = true - return queue.Write(ctx, ev) + executeFn: func(ctx context.Context, reqCtx *a2asrv.ExecutorContext) iter.Seq2[a2a.Event, error] { + return func(yield func(a2a.Event, error) bool) { + remoteTaskIDChan <- reqCtx.TaskID + if !yield(a2a.NewSubmittedTask(reqCtx, reqCtx.Message), nil) { + return + } + ev := a2a.NewStatusUpdateEvent(reqCtx, a2a.TaskStateInputRequired, a2a.NewMessage(a2a.MessageRoleAgent, &a2a.Part{ + Content: a2a.Data{Value: map[string]any{"id": "call-1", "name": "foo"}}, + Metadata: map[string]any{"adk_is_long_running": true, "adk_type": "function_call"}, + })) + yield(ev, nil) + } }, - cancelFn: func(ctx context.Context, reqCtx *a2asrv.RequestContext, queue eventqueue.Queue) error { - ev := a2a.NewStatusUpdateEvent(reqCtx, a2a.TaskStateCanceled, nil) - ev.Final = true - return queue.Write(ctx, ev) + cancelFn: func(ctx context.Context, reqCtx *a2asrv.ExecutorContext) iter.Seq2[a2a.Event, error] { + return func(yield func(a2a.Event, error) bool) { + yield(a2a.NewStatusUpdateEvent(reqCtx, a2a.TaskStateCanceled, nil), nil) + } }, }) defer serverB.Close() @@ -1155,14 +1201,14 @@ func TestA2AMultiHopInputRequiredCancellation(t *testing.T) { // Send message clientA := newA2AClient(t, serverA) - msg1 := a2a.NewMessage(a2a.MessageRoleUser, a2a.TextPart{Text: "Hello"}) + msg1 := a2a.NewMessage(a2a.MessageRoleUser, a2a.NewTextPart("Hello")) task1 := mustSendMessage(t, clientA, msg1) if task1.Status.State != a2a.TaskStateInputRequired { t.Fatalf("task1.Status.State = %q, want %q", task1.Status.State, a2a.TaskStateInputRequired) } // Cancel the task on Server A - _, err := clientA.CancelTask(t.Context(), &a2a.TaskIDParams{ID: task1.ID}) + _, err := clientA.CancelTask(t.Context(), &a2a.CancelTaskRequest{ID: task1.ID}) if err != nil { t.Fatalf("client.CancelTask() error = %v", err) } @@ -1170,7 +1216,7 @@ func TestA2AMultiHopInputRequiredCancellation(t *testing.T) { // Verify that Server B's task was cancelled remoteTaskID := <-remoteTaskIDChan clientB := newA2AClient(t, serverB) - remoteTask, err := clientB.GetTask(t.Context(), &a2a.TaskQueryParams{ID: remoteTaskID}) + remoteTask, err := clientB.GetTask(t.Context(), &a2a.GetTaskRequest{ID: remoteTaskID}) if err != nil { t.Fatalf("client.CancelTask() error = %v", err) } @@ -1178,3 +1224,53 @@ func TestA2AMultiHopInputRequiredCancellation(t *testing.T) { t.Fatalf("remoteTask.Status.State = %q, want %q", remoteTask.Status.State, a2a.TaskStateCanceled) } } + +func TestA2AMultiHopStructuredErrorPropagation(t *testing.T) { + // Server B with structured error serialization logic + structuredErr := a2a.NewDataPart(map[string]any{"error_type": "auth_required"}) + serverB := startA2AServer(adka2a.NewExecutor(adka2a.ExecutorConfig{ + RunnerConfig: runner.Config{ + AppName: "broken", + Agent: utils.Must(agent.New(agent.Config{ + Name: "broken", + Run: func(ic agent.InvocationContext) iter.Seq2[*session.Event, error] { + return func(yield func(*session.Event, error) bool) { + yield(nil, a2a.ErrUnauthorized) + } + }, + })), + SessionService: session.InMemoryService(), + }, + AfterExecuteCallback: func(ctx adka2a.ExecutorContext, finalEvent *a2a.TaskStatusUpdateEvent, err error) error { + if errors.Is(err, a2a.ErrUnauthorized) { + finalEvent.Status.Message.Parts = append(finalEvent.Status.Message.Parts, structuredErr) + } + return nil + }, + })) + defer serverB.Close() + + // Server A, default configuration + remoteAgent := newA2ARemoteAgent(t, "broken-remote", serverB) + rootAgent := newRootAgent("root", remoteAgent) + executorA := newAgentExecutor(rootAgent, nil, adka2a.OutputArtifactPerRun) + serverA := startA2AServer(executorA) + defer serverA.Close() + + // Send message + clientA := newA2AClient(t, serverA) + msg1 := a2a.NewMessage(a2a.MessageRoleUser, a2a.NewTextPart("Hello")) + task1 := mustSendMessage(t, clientA, msg1) + if task1.Status.State != a2a.TaskStateFailed { + t.Fatalf("task1.Status.State = %q, want %q", task1.Status.State, a2a.TaskStateFailed) + } + if len(task1.Status.Message.Parts) != 2 { + t.Fatalf("len(task1.Status.Message.Parts) = %d, want len([error text, extra parts]) = 2", len(task1.Status.Message.Parts)) + } + if !strings.Contains(task1.Status.Message.Parts[0].Text(), a2a.ErrUnauthorized.Error()) { + t.Fatalf("status.Message.Parts[0].Text() = %s, want contain %q", task1.Status.Message.Parts[0].Text(), a2a.ErrUnauthorized.Error()) + } + if diff := cmp.Diff(task1.Status.Message.Parts[1].Data(), structuredErr.Data()); diff != "" { + t.Fatalf("wrong structured error part (-want,+got):\n%s", diff) + } +} diff --git a/agent/remoteagent/v2/client.go b/agent/remoteagent/v2/client.go new file mode 100644 index 000000000..bd445cca0 --- /dev/null +++ b/agent/remoteagent/v2/client.go @@ -0,0 +1,52 @@ +// Copyright 2025 Google LLC +// +// Licensed under the Apache License, Version 2.0 (the "License"); +// you may not use this file except in compliance with the License. +// You may obtain a copy of the License at +// +// http://www.apache.org/licenses/LICENSE-2.0 +// +// Unless required by applicable law or agreed to in writing, software +// distributed under the License is distributed on an "AS IS" BASIS, +// WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. +// See the License for the specific language governing permissions and +// limitations under the License. + +package remoteagent + +import ( + "context" + + "github.com/a2aproject/a2a-go/v2/a2a" + "github.com/a2aproject/a2a-go/v2/a2aclient" + + iremoteagent "google.golang.org/adk/internal/agent/remoteagent" +) + +// A2AClient is used to send messages to a remote agent. +type A2AClient iremoteagent.A2AClient + +// A2AClientProvider is a function that creates an A2AMessageSender. +type A2AClientProvider func(context.Context, *a2a.AgentCard) (A2AClient, error) + +// CreateClient implements iremoteagent.A2AClientProvider. +func (fn A2AClientProvider) CreateClient(ctx context.Context, card *a2a.AgentCard) (iremoteagent.A2AClient, error) { + return fn(ctx, card) +} + +// NewA2AClientProvider creates a default A2AClientProvider from the configured factory. +func NewA2AClientProvider(factory *a2aclient.Factory) A2AClientProvider { + return func(ctx context.Context, card *a2a.AgentCard) (A2AClient, error) { + var client *a2aclient.Client + var err error + if factory != nil { + client, err = factory.CreateFromCard(ctx, card) + } else { + client, err = a2aclient.NewFromCard(ctx, card) + } + if err != nil { + return nil, err + } + return client, nil + } +} diff --git a/agent/remoteagent/doc.go b/agent/remoteagent/v2/doc.go similarity index 91% rename from agent/remoteagent/doc.go rename to agent/remoteagent/v2/doc.go index c26fba2a4..98e2ce114 100644 --- a/agent/remoteagent/doc.go +++ b/agent/remoteagent/v2/doc.go @@ -12,5 +12,5 @@ // See the License for the specific language governing permissions and // limitations under the License. -// Package remoteagent allows to use a remote ADK agents. +// Package remoteagent allows using a remote ADK agents. package remoteagent diff --git a/agent/remoteagent/testdata/TestA2ARemoteAgentStreamingGeminiError.httprr b/agent/remoteagent/v2/testdata/TestA2ARemoteAgentStreamingGeminiError.httprr similarity index 100% rename from agent/remoteagent/testdata/TestA2ARemoteAgentStreamingGeminiError.httprr rename to agent/remoteagent/v2/testdata/TestA2ARemoteAgentStreamingGeminiError.httprr diff --git a/agent/remoteagent/testdata/TestA2ARemoteAgentStreamingGeminiSuccess.httprr b/agent/remoteagent/v2/testdata/TestA2ARemoteAgentStreamingGeminiSuccess.httprr similarity index 100% rename from agent/remoteagent/testdata/TestA2ARemoteAgentStreamingGeminiSuccess.httprr rename to agent/remoteagent/v2/testdata/TestA2ARemoteAgentStreamingGeminiSuccess.httprr diff --git a/agent/remoteagent/testdata/TestA2ASingleHopFinalResponse_llm_mid-response_error.httprr b/agent/remoteagent/v2/testdata/TestA2ASingleHopFinalResponse_llm_mid-response_error.httprr similarity index 100% rename from agent/remoteagent/testdata/TestA2ASingleHopFinalResponse_llm_mid-response_error.httprr rename to agent/remoteagent/v2/testdata/TestA2ASingleHopFinalResponse_llm_mid-response_error.httprr diff --git a/agent/remoteagent/testdata/TestA2ASingleHopFinalResponse_llm_mid-response_error_response.httprr b/agent/remoteagent/v2/testdata/TestA2ASingleHopFinalResponse_llm_mid-response_error_response.httprr similarity index 100% rename from agent/remoteagent/testdata/TestA2ASingleHopFinalResponse_llm_mid-response_error_response.httprr rename to agent/remoteagent/v2/testdata/TestA2ASingleHopFinalResponse_llm_mid-response_error_response.httprr diff --git a/agent/remoteagent/utils.go b/agent/remoteagent/v2/utils.go similarity index 83% rename from agent/remoteagent/utils.go rename to agent/remoteagent/v2/utils.go index bca2e4c8c..ae02b157b 100644 --- a/agent/remoteagent/utils.go +++ b/agent/remoteagent/v2/utils.go @@ -18,11 +18,12 @@ import ( "fmt" "slices" - "github.com/a2aproject/a2a-go/a2a" + "github.com/a2aproject/a2a-go/v2/a2a" + "github.com/a2aproject/a2a-go/v2/log" "google.golang.org/genai" "google.golang.org/adk/agent" - "google.golang.org/adk/server/adka2a" + "google.golang.org/adk/server/adka2a/v2" "google.golang.org/adk/session" ) @@ -89,23 +90,23 @@ func getFunctionResponseCallID(event *session.Event) (string, bool) { // Parts from all events we processed are returned as a single list. // The returned contextID might be an empty string. This means the current remote agent invocation is not associates with // any of the previous one. In this case a new contextID will be generated on the remote server. -func toMissingRemoteSessionParts(ctx agent.InvocationContext, events session.Events, cfg A2AConfig) ([]a2a.Part, string) { +func toMissingRemoteSessionParts(ctx agent.InvocationContext, events session.Events, cfg A2AConfig) ([]*a2a.Part, string) { partCount, contextID := 0, "" // only events after this index are not in the remote session lastRemoteResponseIndex := -1 for i := events.Len() - 1; i >= 0; i-- { event := events.At(i) - if event.LLMResponse.Content != nil { - partCount += len(event.Content.Parts) - } if event.Author == ctx.Agent().Name() { lastRemoteResponseIndex = i _, contextID = adka2a.GetA2ATaskInfo(event) break } + if event.LLMResponse.Content != nil { + partCount += len(event.Content.Parts) + } } - result := make([]a2a.Part, 0, partCount) + result := make([]*a2a.Part, 0, partCount) for i := lastRemoteResponseIndex + 1; i < events.Len(); i++ { event := events.At(i) if event.Author != "user" && event.Author != ctx.Agent().Name() { @@ -116,7 +117,7 @@ func toMissingRemoteSessionParts(ctx agent.InvocationContext, events session.Eve } parts, err := convertParts(ctx, cfg, event) if err != nil { - // TODO(yarolegovich): log error + log.Warn(ctx, "failed to convert parts for session event", "index", i, "error", err) continue } result = append(result, parts...) @@ -125,7 +126,7 @@ func toMissingRemoteSessionParts(ctx agent.InvocationContext, events session.Eve } func presentAsUserMessage(ctx agent.InvocationContext, agentEvent *session.Event) *session.Event { - event := session.NewEvent(ctx.InvocationID()) + event := session.NewEventWithContext(ctx, ctx.InvocationID()) event.Author = "user" if agentEvent.Content == nil { @@ -158,3 +159,25 @@ func presentAsUserMessage(ctx agent.InvocationContext, agentEvent *session.Event } return event } + +func convertParts(ctx agent.InvocationContext, cfg A2AConfig, event *session.Event) ([]*a2a.Part, error) { + parts := make([]*a2a.Part, 0, len(event.Content.Parts)) + if cfg.GenAIPartConverter != nil { + for _, part := range event.Content.Parts { + cp, err := cfg.GenAIPartConverter(ctx, event, part) + if err != nil { + return nil, err + } + if cp != nil { + parts = append(parts, cp) + } + } + } else { + var err error + parts, err = adka2a.ToA2AParts(event.Content.Parts, event.LongRunningToolIDs) + if err != nil { + return nil, fmt.Errorf("event part conversion failed: %w", err) + } + } + return parts, nil +} diff --git a/agent/remoteagent/utils_test.go b/agent/remoteagent/v2/utils_test.go similarity index 94% rename from agent/remoteagent/utils_test.go rename to agent/remoteagent/v2/utils_test.go index d189c4e29..1690f7432 100644 --- a/agent/remoteagent/utils_test.go +++ b/agent/remoteagent/v2/utils_test.go @@ -18,7 +18,7 @@ import ( "fmt" "testing" - "github.com/a2aproject/a2a-go/a2a" + "github.com/a2aproject/a2a-go/v2/a2a" "github.com/google/go-cmp/cmp" "github.com/google/go-cmp/cmp/cmpopts" "google.golang.org/genai" @@ -26,7 +26,7 @@ import ( "google.golang.org/adk/agent" icontext "google.golang.org/adk/internal/context" "google.golang.org/adk/model" - "google.golang.org/adk/server/adka2a" + "google.golang.org/adk/server/adka2a/v2" "google.golang.org/adk/session" ) @@ -157,7 +157,7 @@ func TestToMissingRemoteSessionParts(t *testing.T) { testCases := []struct { name string events []*session.Event - wantParts []a2a.Part + wantParts []*a2a.Part wantContextID string }{ { @@ -166,10 +166,10 @@ func TestToMissingRemoteSessionParts(t *testing.T) { newEventFromParts("user", &genai.Part{Text: "hello"}), newEventFromParts("user", &genai.Part{Text: "foo"}, &genai.Part{Text: "bar"}), }, - wantParts: []a2a.Part{ - a2a.TextPart{Text: "hello"}, - a2a.TextPart{Text: "foo"}, - a2a.TextPart{Text: "bar"}, + wantParts: []*a2a.Part{ + a2a.NewTextPart("hello"), + a2a.NewTextPart("foo"), + a2a.NewTextPart("bar"), }, }, { @@ -178,10 +178,10 @@ func TestToMissingRemoteSessionParts(t *testing.T) { newEventFromParts("another-agent", &genai.Part{Text: "foo"}), newEventFromParts("user", &genai.Part{Text: "bar"}), }, - wantParts: []a2a.Part{ - a2a.TextPart{Text: "For context:"}, - a2a.TextPart{Text: "[another-agent] said: foo"}, - a2a.TextPart{Text: "bar"}, + wantParts: []*a2a.Part{ + a2a.NewTextPart("For context:"), + a2a.NewTextPart("[another-agent] said: foo"), + a2a.NewTextPart("bar"), }, }, { @@ -190,8 +190,8 @@ func TestToMissingRemoteSessionParts(t *testing.T) { newEventFromParts("another-agent", &genai.Part{Text: "foo", Thought: true}), newEventFromParts("user", &genai.Part{Text: "bar"}), }, - wantParts: []a2a.Part{ - a2a.TextPart{Text: "bar"}, + wantParts: []*a2a.Part{ + a2a.NewTextPart("bar"), }, }, { @@ -202,9 +202,9 @@ func TestToMissingRemoteSessionParts(t *testing.T) { newEventFromParts("user", &genai.Part{Text: "foo"}), newEventFromParts("user", &genai.Part{Text: "bar"}), }, - wantParts: []a2a.Part{ - a2a.TextPart{Text: "foo"}, - a2a.TextPart{Text: "bar"}, + wantParts: []*a2a.Part{ + a2a.NewTextPart("foo"), + a2a.NewTextPart("bar"), }, }, { @@ -218,7 +218,7 @@ func TestToMissingRemoteSessionParts(t *testing.T) { }, }, }, - wantParts: []a2a.Part{}, + wantParts: []*a2a.Part{}, wantContextID: "ctxID-123", }, } diff --git a/agent/workflowagents/loopagent/agent_test.go b/agent/workflowagents/loopagent/agent_test.go index 690b12418..d21f361da 100644 --- a/agent/workflowagents/loopagent/agent_test.go +++ b/agent/workflowagents/loopagent/agent_test.go @@ -306,13 +306,13 @@ func (a *customAgent) Run(agent.InvocationContext) iter.Seq2[*session.Event, err type EmptyArgs struct{} -func exampleFunctionThatEscalates(ctx tool.Context, myArgs EmptyArgs) (map[string]string, error) { +func exampleFunctionThatEscalates(ctx agent.ToolContext, myArgs EmptyArgs) (map[string]string, error) { ctx.Actions().Escalate = true ctx.Actions().SkipSummarization = false return map[string]string{}, nil } -func exampleFunctionThatEscalatesAndSkips(ctx tool.Context, myArgs EmptyArgs) (map[string]string, error) { +func exampleFunctionThatEscalatesAndSkips(ctx agent.ToolContext, myArgs EmptyArgs) (map[string]string, error) { ctx.Actions().Escalate = true ctx.Actions().SkipSummarization = true return map[string]string{}, nil diff --git a/agent/workflowagents/parallelagent/agent_test.go b/agent/workflowagents/parallelagent/agent_test.go index bda8f709d..1617ad2e0 100644 --- a/agent/workflowagents/parallelagent/agent_test.go +++ b/agent/workflowagents/parallelagent/agent_test.go @@ -298,7 +298,7 @@ func createAgentWithGemini(t *testing.T, name string) agent.Agent { Name: fmt.Sprintf("search_tool_%s", name), Description: "Search for information on the web", }, - func(ctx tool.Context, args struct{ Query string }) (string, error) { + func(ctx agent.ToolContext, args struct{ Query string }) (string, error) { return fmt.Sprintf("search result for '%s' from %s", args.Query, name), nil }, ) @@ -311,7 +311,7 @@ func createAgentWithGemini(t *testing.T, name string) agent.Agent { Name: fmt.Sprintf("analyze_tool_%s", name), Description: "Analyze data and return insights", }, - func(ctx tool.Context, args struct{ Data string }) (string, error) { + func(ctx agent.ToolContext, args struct{ Data string }) (string, error) { return fmt.Sprintf("analysis result for '%s' from %s", args.Data, name), nil }, ) diff --git a/agent/workflowagents/sequentialagent/agent.go b/agent/workflowagents/sequentialagent/agent.go index b307aa6b3..0c648175e 100644 --- a/agent/workflowagents/sequentialagent/agent.go +++ b/agent/workflowagents/sequentialagent/agent.go @@ -18,16 +18,29 @@ package sequentialagent import ( "fmt" "iter" + "log" + "sync" "google.golang.org/adk/agent" agentinternal "google.golang.org/adk/internal/agent" + "google.golang.org/adk/internal/llminternal" "google.golang.org/adk/session" + "google.golang.org/adk/tool/functiontool" ) // New creates a SequentialAgent. // // SequentialAgent executes its sub-agents once, in the order they are listed. -// +type seqAgent struct { + agent.Agent + *agentinternal.State + impl *sequentialAgent +} + +func (s *seqAgent) RunLive(ctx agent.InvocationContext) (agent.LiveSession, iter.Seq2[*session.Event, error], error) { + return s.impl.RunLive(ctx) +} + // Use the SequentialAgent when you want the execution to occur in a fixed, // strict order. func New(cfg Config) (agent.Agent, error) { @@ -51,7 +64,7 @@ func New(cfg Config) (agent.Agent, error) { state.AgentType = agentinternal.TypeSequentialAgent state.Config = cfg - return sequentialAgent, nil + return &seqAgent{Agent: sequentialAgent, State: state, impl: sequentialAgentImpl}, nil } // Config defines the configuration for a SequentialAgent. @@ -74,3 +87,118 @@ func (a *sequentialAgent) Run(ctx agent.InvocationContext) iter.Seq2[*session.Ev } } } + +type sequentialLiveSession struct { + mu sync.Mutex + activeSess agent.LiveSession + closed bool +} + +func (s *sequentialLiveSession) Send(req agent.LiveRequest) error { + s.mu.Lock() + defer s.mu.Unlock() + if s.closed { + return fmt.Errorf("session is closed") + } + if s.activeSess == nil { + return fmt.Errorf("no active sub-agent live session") + } + return s.activeSess.Send(req) +} + +func (s *sequentialLiveSession) Close() error { + s.mu.Lock() + defer s.mu.Unlock() + s.closed = true + if s.activeSess != nil { + return s.activeSess.Close() + } + return nil +} + +func (s *sequentialLiveSession) setActiveSession(sess agent.LiveSession) { + s.mu.Lock() + defer s.mu.Unlock() + s.activeSess = sess +} + +func (a *sequentialAgent) RunLive(ctx agent.InvocationContext) (agent.LiveSession, iter.Seq2[*session.Event, error], error) { + subAgents := ctx.Agent().SubAgents() + if len(subAgents) == 0 { + return nil, nil, fmt.Errorf("sequential agent has no sub-agents") + } + + // Inject task_completed tool into sub LLM agents + type taskCompletedArgs struct{} + type taskCompletedResults struct { + Result string `json:"result"` + } + + taskCompletedTool, err := functiontool.New(functiontool.Config{ + Name: "task_completed", + Description: "Signals that the agent has successfully completed the user's question or task.", + }, func(ctx agent.ToolContext, args taskCompletedArgs) (taskCompletedResults, error) { + return taskCompletedResults{Result: "Task completion signaled."}, nil + }) + if err != nil { + return nil, nil, fmt.Errorf("failed to create task_completed tool: %w", err) + } + + for _, subAgent := range subAgents { + if llmAgent, ok := subAgent.(llminternal.Agent); ok { + state := llminternal.Reveal(llmAgent) + hasTaskCompleted := false + for _, t := range state.Tools { + if t.Name() == "task_completed" { + hasTaskCompleted = true + break + } + } + if !hasTaskCompleted { + state.Tools = append(state.Tools, taskCompletedTool) + instructionSuffix := "\nIf you finished the user's request according to its description, call the task_completed function to exit so the next agents can take over. When calling this function, do not generate any text other than the function call." + state.Instruction += instructionSuffix + } + } + } + + seqSess := &sequentialLiveSession{} + + wrappedIter := func(yield func(*session.Event, error) bool) { + for _, subAgent := range subAgents { + liveAgent, ok := subAgent.(interface { + RunLive(ctx agent.InvocationContext) (agent.LiveSession, iter.Seq2[*session.Event, error], error) + }) + if !ok { + if !yield(nil, fmt.Errorf("sub-agent %s does not support Live Run", subAgent.Name())) { + return + } + return + } + + subSess, innerIter, err := liveAgent.RunLive(ctx) + if err != nil { + if !yield(nil, fmt.Errorf("sub-agent %s RunLive failed: %w", subAgent.Name(), err)) { + return + } + return + } + + seqSess.setActiveSession(subSess) + + for ev, err := range innerIter { + if !yield(ev, err) { + if err := subSess.Close(); err != nil { + log.Printf("error closing sub-session: %v\n", err) + } + return + } + } + if err := subSess.Close(); err != nil { + log.Printf("error closing sub-session: %v\n", err) + } + } + } + + return seqSess, wrappedIter, nil +} diff --git a/agent/workflowagents/sequentialagent/agent_test.go b/agent/workflowagents/sequentialagent/agent_test.go index 59a805535..a4a9cf7d7 100644 --- a/agent/workflowagents/sequentialagent/agent_test.go +++ b/agent/workflowagents/sequentialagent/agent_test.go @@ -19,6 +19,7 @@ import ( "fmt" "iter" "testing" + "time" "github.com/google/go-cmp/cmp" "github.com/google/go-cmp/cmp/cmpopts" @@ -27,6 +28,7 @@ import ( "google.golang.org/adk/agent" "google.golang.org/adk/agent/llmagent" "google.golang.org/adk/agent/workflowagents/sequentialagent" + "google.golang.org/adk/internal/llminternal" "google.golang.org/adk/model" "google.golang.org/adk/runner" "google.golang.org/adk/session" @@ -335,3 +337,237 @@ func (f *FakeLLM) GenerateContent(ctx context.Context, req *model.LLMRequest, st }, nil) } } + +type mockLiveAgent struct { + agent.Agent + runLiveFn func(ctx agent.InvocationContext) (agent.LiveSession, iter.Seq2[*session.Event, error], error) +} + +func (m *mockLiveAgent) RunLive(ctx agent.InvocationContext) (agent.LiveSession, iter.Seq2[*session.Event, error], error) { + return m.runLiveFn(ctx) +} + +type dummyLiveSession struct { + sendChan chan agent.LiveRequest + closed bool +} + +func (d *dummyLiveSession) Send(req agent.LiveRequest) error { + d.sendChan <- req + return nil +} + +func (d *dummyLiveSession) Close() error { + d.closed = true + return nil +} + +func mustAgent(a agent.Agent, err error) agent.Agent { + if err != nil { + panic(err) + } + return a +} + +type mockInvocationContext struct { + agent.InvocationContext + agent agent.Agent + invocationID string + ctx context.Context +} + +func (m *mockInvocationContext) Agent() agent.Agent { + return m.agent +} + +func (m *mockInvocationContext) InvocationID() string { + return m.invocationID +} + +func (m *mockInvocationContext) Context() context.Context { + return m.ctx +} + +func (m *mockInvocationContext) Deadline() (time.Time, bool) { return m.ctx.Deadline() } +func (m *mockInvocationContext) Done() <-chan struct{} { return m.ctx.Done() } +func (m *mockInvocationContext) Err() error { return m.ctx.Err() } +func (m *mockInvocationContext) Value(key any) any { return m.ctx.Value(key) } + +func TestSequentialAgent_RunLive_Injection(t *testing.T) { + subAgent1 := newCustomAgent(t, 1) + subAgent2 := newCustomAgent(t, 2) + + sequentialAgent, err := sequentialagent.New(sequentialagent.Config{ + AgentConfig: agent.Config{ + Name: "seq_agent", + SubAgents: []agent.Agent{subAgent1, subAgent2}, + }, + }) + if err != nil { + t.Fatalf("failed to create sequential agent: %v", err) + } + + // Before RunLive, sub-agents do not have the task_completed tool + if llmAgent1, ok := subAgent1.(llminternal.Agent); ok { + state := llminternal.Reveal(llmAgent1) + for _, tool := range state.Tools { + if tool.Name() == "task_completed" { + t.Errorf("sub-agent 1 already has task_completed tool before RunLive") + } + } + } + + // Call RunLive (it will prepare/inject but will fail/return error when executing due to nil/mock context, + // which is perfectly fine since the injection happens beforehand). + // Let's pass a mock context that returns seqAgent as the Agent. + invCtx := &mockInvocationContext{ + agent: sequentialAgent, + invocationID: "test_id", + ctx: t.Context(), + } + + liveAgent, ok := sequentialAgent.(interface { + RunLive(ctx agent.InvocationContext) (agent.LiveSession, iter.Seq2[*session.Event, error], error) + }) + if !ok { + t.Fatalf("sequential agent does not implement RunLive") + } + + _, _, _ = liveAgent.RunLive(invCtx) + + // After RunLive initiation, the sub-agents MUST have the task_completed tool injected! + if llmAgent1, ok := subAgent1.(llminternal.Agent); ok { + state := llminternal.Reveal(llmAgent1) + hasTaskCompleted := false + for _, tool := range state.Tools { + if tool.Name() == "task_completed" { + hasTaskCompleted = true + break + } + } + if !hasTaskCompleted { + t.Errorf("sub-agent 1 does not have task_completed tool injected after RunLive") + } + } +} + +func TestSequentialAgent_RunLive_SequentialOrchestration(t *testing.T) { + ctx := t.Context() + + sendChan1 := make(chan agent.LiveRequest, 10) + sendChan2 := make(chan agent.LiveRequest, 10) + + subSess1 := &dummyLiveSession{sendChan: sendChan1} + subSess2 := &dummyLiveSession{sendChan: sendChan2} + + agent1 := mustAgent(agent.New(agent.Config{Name: "sub_agent_1"})) + liveAgent1 := &mockLiveAgent{ + Agent: agent1, + runLiveFn: func(ctx agent.InvocationContext) (agent.LiveSession, iter.Seq2[*session.Event, error], error) { + iterFn := func(yield func(*session.Event, error) bool) { + ev := session.NewEventWithContext(ctx, ctx.InvocationID()) + ev.Author = "sub_agent_1" + yield(ev, nil) + } + return subSess1, iterFn, nil + }, + } + + agent2 := mustAgent(agent.New(agent.Config{Name: "sub_agent_2"})) + liveAgent2 := &mockLiveAgent{ + Agent: agent2, + runLiveFn: func(ctx agent.InvocationContext) (agent.LiveSession, iter.Seq2[*session.Event, error], error) { + iterFn := func(yield func(*session.Event, error) bool) { + ev := session.NewEventWithContext(ctx, ctx.InvocationID()) + ev.Author = "sub_agent_2" + yield(ev, nil) + } + return subSess2, iterFn, nil + }, + } + + seqAgent, err := sequentialagent.New(sequentialagent.Config{ + AgentConfig: agent.Config{ + Name: "seq_agent", + SubAgents: []agent.Agent{liveAgent1, liveAgent2}, + }, + }) + if err != nil { + t.Fatalf("failed to create sequential agent: %v", err) + } + + invCtx := &mockInvocationContext{ + agent: seqAgent, + invocationID: "test_inv_id", + ctx: ctx, + } + + liveAgent, ok := seqAgent.(interface { + RunLive(ctx agent.InvocationContext) (agent.LiveSession, iter.Seq2[*session.Event, error], error) + }) + if !ok { + t.Fatalf("sequential agent does not implement RunLive") + } + + sess, seqIter, err := liveAgent.RunLive(invCtx) + if err != nil { + t.Fatalf("RunLive failed: %v", err) + } + + next, stop := iter.Pull2(seqIter) + defer stop() + + // Consume first sub-agent event + ev1, err1, ok := next() + if !ok || err1 != nil { + t.Fatalf("expected first event, got ok=%v, err=%v", ok, err1) + } + if ev1.Author != "sub_agent_1" { + t.Errorf("expected event from sub_agent_1, got %s", ev1.Author) + } + + // Now seqSess should route to subSess1 + req1 := agent.LiveRequest{Content: genai.NewContentFromText("to agent 1", "")} + if err := sess.Send(req1); err != nil { + t.Fatalf("failed to Send to sess: %v", err) + } + gotReq1 := <-sendChan1 + if gotReq1.Content.Parts[0].Text != "to agent 1" { + t.Errorf("expected request to subSess1, got: %v", gotReq1) + } + + // The subSess1 completes, transitioning to agent2 + ev2, err2, ok := next() + if !ok || err2 != nil { + t.Fatalf("expected second event, got ok=%v, err=%v", ok, err2) + } + if ev2.Author != "sub_agent_2" { + t.Errorf("expected event from sub_agent_2, got %s", ev2.Author) + } + + // Now seqSess should route to subSess2 + req2 := agent.LiveRequest{Content: genai.NewContentFromText("to agent 2", "")} + if err := sess.Send(req2); err != nil { + t.Fatalf("failed to Send to sess: %v", err) + } + gotReq2 := <-sendChan2 + if gotReq2.Content.Parts[0].Text != "to agent 2" { + t.Errorf("expected request to subSess2, got: %v", gotReq2) + } + + // Verify that subSess1 is closed + if !subSess1.closed { + t.Errorf("expected sub_agent_1 session to be closed after transition") + } + + // The subSess2 completes + _, _, ok = next() + if ok { + t.Errorf("expected iterator to be exhausted") + } + + // Verify subSess2 is closed + if !subSess2.closed { + t.Errorf("expected sub_agent_2 session to be closed at the end") + } +} diff --git a/cmd/adkgo/internal/deploy/agentengine/agentengine.go b/cmd/adkgo/internal/deploy/agentengine/agentengine.go index 206fa16ca..2e03f9e2f 100644 --- a/cmd/adkgo/internal/deploy/agentengine/agentengine.go +++ b/cmd/adkgo/internal/deploy/agentengine/agentengine.go @@ -28,10 +28,12 @@ import ( "strings" "time" - aiplatform "cloud.google.com/go/aiplatform/apiv1" - "cloud.google.com/go/aiplatform/apiv1/aiplatformpb" + aiplatform "cloud.google.com/go/aiplatform/apiv1beta1" + "cloud.google.com/go/aiplatform/apiv1beta1/aiplatformpb" "github.com/spf13/cobra" "google.golang.org/api/option" + "google.golang.org/protobuf/types/known/durationpb" + "google.golang.org/protobuf/types/known/fieldmaskpb" "google.golang.org/adk/cmd/adkgo/internal/deploy" "google.golang.org/adk/internal/cli/util" @@ -43,10 +45,18 @@ type gCloudFlags struct { projectName string } +type memoryBankFlags struct { + deploy bool + model string + ttl time.Duration +} + type agentEngineServiceFlags struct { - name string - displayName string - serverPort int + name string + displayName string + serverPort int + agentEngineID string + memoryBank memoryBankFlags } type buildFlags struct { @@ -79,7 +89,7 @@ var agentEngineCmd = &cobra.Command{ Short: "Deploys the application to Agent Engine.", Long: `Deploys the application to Agent Engine. It creates a source archive, uploads it to create a Reasoning Engine, and cleans up temporary files.`, RunE: func(cmd *cobra.Command, args []string) error { - return flags.deployOnagentEngine() + return flags.deployOnAgentEngine() }, } @@ -94,6 +104,12 @@ func init() { agentEngineCmd.PersistentFlags().IntVar(&flags.agentEngine.serverPort, "server_port", 8080, "agentEngine server port") agentEngineCmd.PersistentFlags().StringVarP(&flags.source.entryPointPath, "entry_point_path", "e", "", "Path to an entry point (go 'main')") agentEngineCmd.PersistentFlags().StringVarP(&flags.source.sourceDir, "source_dir", "d", "", "Directory to archive, defaults to current working directory") + agentEngineCmd.PersistentFlags().StringVar(&flags.agentEngine.agentEngineID, "agent_engine_id", "", "ID of the Agent Engine instance to update if it exists (default: \"\", which means a new instance will be created).") + agentEngineCmd.PersistentFlags().BoolVar(&flags.agentEngine.memoryBank.deploy, "mem_deploy", false, "If set to true then memory bank will be deployed too") + agentEngineCmd.PersistentFlags().StringVar(&flags.agentEngine.memoryBank.model, "mem_model", "publishers/google/models/gemini-2.5-flash", "Name of the model to be used for memory generation - for list you can GET"+ + " https://${LOCATION_ID}-aiplatform.googleapis.com/v1beta1/publishers/google/models. It will be prefixed with projects//locations/ forming for instance full model name like "+ + " 'projects/project_name/locations/us-central1/publishers/google/models/gemini-2.5-flash'") + agentEngineCmd.PersistentFlags().DurationVar(&flags.agentEngine.memoryBank.ttl, "mem_ttl", time.Hour*24*365, "Time-To-Live for memories") } // computeFlags uses command line arguments to create a full config @@ -141,6 +157,9 @@ func (f *deployAgentEngineFlags) computeFlags() error { f.agentEngine.displayName = "ADK Agent: " + dateTimeString } + // add provided model as suffix + f.agentEngine.memoryBank.model = fmt.Sprintf("projects/%s/locations/%s/%s", f.gcloud.projectName, f.gcloud.region, f.agentEngine.memoryBank.model) + return nil }) } @@ -269,6 +288,22 @@ func (f *deployAgentEngineFlags) gcloudDeployToAgentEngine() error { }, }, } + + if f.agentEngine.memoryBank.deploy { + req.ReasoningEngine.ContextSpec = &aiplatformpb.ReasoningEngineContextSpec{ + MemoryBankConfig: &aiplatformpb.ReasoningEngineContextSpec_MemoryBankConfig{ + GenerationConfig: &aiplatformpb.ReasoningEngineContextSpec_MemoryBankConfig_GenerationConfig{ + Model: f.agentEngine.memoryBank.model, + }, + TtlConfig: &aiplatformpb.ReasoningEngineContextSpec_MemoryBankConfig_TtlConfig{ + Ttl: &aiplatformpb.ReasoningEngineContextSpec_MemoryBankConfig_TtlConfig_DefaultTtl{ + DefaultTtl: durationpb.New(f.agentEngine.memoryBank.ttl), + }, + }, + }, + } + } + p("Sending CreateReasoningEngine request...") op, err := client.CreateReasoningEngine(ctx, req) if err != nil { @@ -288,8 +323,97 @@ func (f *deployAgentEngineFlags) gcloudDeployToAgentEngine() error { }) } -// deployOnagentEngine executes the sequence of actions preparing and deploying the agent to agentEngine -func (f *deployAgentEngineFlags) deployOnagentEngine() error { +// gcloudUpdateAgentEngine invokes gcloud to update source on agentEngine +func (f *deployAgentEngineFlags) gcloudUpdateAgentEngine() error { + return util.LogStartStop("Updating Agent Engine", + func(p util.Printer) error { + ctx := context.Background() + name := fmt.Sprintf("projects/%s/locations/%s/reasoningEngines/%s", f.gcloud.projectName, f.gcloud.region, f.agentEngine.agentEngineID) + endpoint := fmt.Sprintf("%s-aiplatform.googleapis.com:443", f.gcloud.region) + client, err := aiplatform.NewReasoningEngineClient(ctx, option.WithEndpoint(endpoint)) + if err != nil { + return fmt.Errorf("cannot create ReasoningEngineClient: %w", err) + } + defer func() { + if err := client.Close(); err != nil { + p("Warning: failed to close ReasoningEngineClient: %v", err) + } + }() + + archiveContent, err := os.ReadFile(f.build.archivePath) + if err != nil { + return fmt.Errorf("cannot read archive file: %w", err) + } + + methods, err := agentengine.ListClassMethods() + if err != nil { + return fmt.Errorf("cannot list class methods: %w", err) + } + methodsJSON, err := json.Marshal(methods) + if err != nil { + return fmt.Errorf("cannot marshal methods: %w", err) + } + p("Methods:", string(methodsJSON)) + + req := &aiplatformpb.UpdateReasoningEngineRequest{ + ReasoningEngine: &aiplatformpb.ReasoningEngine{ + Name: name, + Spec: &aiplatformpb.ReasoningEngineSpec{ + DeploymentSource: &aiplatformpb.ReasoningEngineSpec_SourceCodeSpec_{ + SourceCodeSpec: &aiplatformpb.ReasoningEngineSpec_SourceCodeSpec{ + Source: &aiplatformpb.ReasoningEngineSpec_SourceCodeSpec_InlineSource_{ + InlineSource: &aiplatformpb.ReasoningEngineSpec_SourceCodeSpec_InlineSource{ + SourceArchive: archiveContent, + }, + }, + LanguageSpec: &aiplatformpb.ReasoningEngineSpec_SourceCodeSpec_ImageSpec_{ + ImageSpec: &aiplatformpb.ReasoningEngineSpec_SourceCodeSpec_ImageSpec{}, + }, + }, + }, + ClassMethods: methods, + }, + }, + UpdateMask: &fieldmaskpb.FieldMask{Paths: []string{"spec.source_code_spec", "spec.class_methods"}}, + } + + if f.agentEngine.memoryBank.deploy { + req.ReasoningEngine.ContextSpec = &aiplatformpb.ReasoningEngineContextSpec{ + MemoryBankConfig: &aiplatformpb.ReasoningEngineContextSpec_MemoryBankConfig{ + GenerationConfig: &aiplatformpb.ReasoningEngineContextSpec_MemoryBankConfig_GenerationConfig{ + Model: f.agentEngine.memoryBank.model, + }, + TtlConfig: &aiplatformpb.ReasoningEngineContextSpec_MemoryBankConfig_TtlConfig{ + Ttl: &aiplatformpb.ReasoningEngineContextSpec_MemoryBankConfig_TtlConfig_DefaultTtl{ + DefaultTtl: durationpb.New(f.agentEngine.memoryBank.ttl), + }, + }, + }, + } + req.UpdateMask.Paths = append(req.UpdateMask.Paths, "spec.context_spec") + } + + p("Sending UpdateReasoningEngine request...") + op, err := client.UpdateReasoningEngine(ctx, req) + if err != nil { + return fmt.Errorf("UpdateReasoningEngine failed: %w", err) + } + + p("Waiting for operation to complete...") + re, err := op.Wait(ctx) + if err != nil { + return fmt.Errorf("operation failed: %w", err) + } + + p("Updated Reasoning Engine:", re.Name) + p("Display Name:", re.DisplayName) + + return nil + }) +} + +// deployOnAgentEngine executes the sequence of actions preparing and deploying the agent to agentEngine +func (f *deployAgentEngineFlags) deployOnAgentEngine() error { fmt.Println(flags) err := f.computeFlags() @@ -304,7 +428,11 @@ func (f *deployAgentEngineFlags) deployOnagentEngine() error { if err != nil { return err } - err = f.gcloudDeployToAgentEngine() + if f.agentEngine.agentEngineID != "" { + err = f.gcloudUpdateAgentEngine() + } else { + err = f.gcloudDeployToAgentEngine() + } if err != nil { return err } diff --git a/cmd/internal/adkcli/main.go b/cmd/internal/adkcli/main.go index 990643dc3..7714042b1 100644 --- a/cmd/internal/adkcli/main.go +++ b/cmd/internal/adkcli/main.go @@ -27,6 +27,7 @@ import ( "google.golang.org/adk/cmd/launcher/full" "google.golang.org/adk/internal/configurable" "google.golang.org/adk/internal/configurable/conformance" + "google.golang.org/adk/internal/configurable/conformance/recordplugin" "google.golang.org/adk/internal/configurable/conformance/replayplugin" "google.golang.org/adk/plugin" "google.golang.org/adk/runner" @@ -113,6 +114,7 @@ func main() { PluginConfig: runner.PluginConfig{ Plugins: []*plugin.Plugin{ replayplugin.MustNew(cwd), + recordplugin.MustNew(cwd), }, }, } diff --git a/cmd/launcher/console/console.go b/cmd/launcher/console/console.go index 5b7f69da1..f40b4e26e 100644 --- a/cmd/launcher/console/console.go +++ b/cmd/launcher/console/console.go @@ -107,6 +107,7 @@ func (l *consoleLauncher) Run(ctx context.Context, config *launcher.Config) erro SessionService: sessionService, ArtifactService: config.ArtifactService, PluginConfig: config.PluginConfig, + MemoryService: config.MemoryService, }) if err != nil { return fmt.Errorf("failed to create runner: %v", err) diff --git a/cmd/launcher/launcher.go b/cmd/launcher/launcher.go index 82570ce22..2642bcaaa 100644 --- a/cmd/launcher/launcher.go +++ b/cmd/launcher/launcher.go @@ -18,7 +18,7 @@ package launcher import ( "context" - "github.com/a2aproject/a2a-go/a2asrv" + "github.com/a2aproject/a2a-go/v2/a2asrv" "google.golang.org/adk/agent" "google.golang.org/adk/artifact" diff --git a/cmd/launcher/web/a2a/a2a.go b/cmd/launcher/web/a2a/a2a.go index 06d7b1fbd..b69558208 100644 --- a/cmd/launcher/web/a2a/a2a.go +++ b/cmd/launcher/web/a2a/a2a.go @@ -20,19 +20,23 @@ import ( "fmt" "net/url" - a2acore "github.com/a2aproject/a2a-go/a2a" - "github.com/a2aproject/a2a-go/a2asrv" + a2acore "github.com/a2aproject/a2a-go/v2/a2a" + "github.com/a2aproject/a2a-go/v2/a2acompat/a2av0" + "github.com/a2aproject/a2a-go/v2/a2asrv" "github.com/gorilla/mux" "google.golang.org/adk/cmd/launcher" "google.golang.org/adk/cmd/launcher/web" "google.golang.org/adk/internal/cli/util" "google.golang.org/adk/runner" - "google.golang.org/adk/server/adka2a" + "google.golang.org/adk/server/adka2a/v2" ) -// apiPath is a suffix used to build an A2A invocation URL -const apiPath = "/a2a/invoke" +// compatAPIPath is a suffix used to build an A2A invocation URL for 0.3 +const compatAPIPath = "/a2a/invoke" + +// apiPath is a suffix used to build an A2A invocation URL for 1.0 +const apiPath = "/a2a/v1/invoke" // a2aConfig contains parameters for launching ADK A2A server type a2aConfig struct { @@ -79,6 +83,10 @@ func (a *a2aLauncher) Parse(args []string) ([]string, error) { // SetupSubrouters implements the web.Sublauncher interface. It adds A2A paths to the main router. func (a *a2aLauncher) SetupSubrouters(router *mux.Router, config *launcher.Config) error { + publicCompatURL, err := url.JoinPath(a.config.agentURL, compatAPIPath) + if err != nil { + return err + } publicURL, err := url.JoinPath(a.config.agentURL, apiPath) if err != nil { return err @@ -86,30 +94,46 @@ func (a *a2aLauncher) SetupSubrouters(router *mux.Router, config *launcher.Confi rootAgent := config.AgentLoader.RootAgent() agentCard := &a2acore.AgentCard{ - Name: rootAgent.Name(), - Description: rootAgent.Description(), - DefaultInputModes: []string{"text/plain"}, - DefaultOutputModes: []string{"text/plain"}, - URL: publicURL, - PreferredTransport: a2acore.TransportProtocolJSONRPC, - Skills: adka2a.BuildAgentSkills(rootAgent), - Capabilities: a2acore.AgentCapabilities{Streaming: true}, - SupportsAuthenticatedExtendedCard: false, + Name: rootAgent.Name(), + Description: rootAgent.Description(), + DefaultInputModes: []string{"text/plain"}, + DefaultOutputModes: []string{"text/plain"}, + SupportedInterfaces: []*a2acore.AgentInterface{ + { + URL: publicURL, + ProtocolBinding: a2acore.TransportProtocolJSONRPC, + ProtocolVersion: a2acore.Version, + }, + { + URL: publicCompatURL, + ProtocolBinding: a2acore.TransportProtocolJSONRPC, + ProtocolVersion: a2av0.Version, + }, + }, + Version: "1.0.0", + Skills: adka2a.BuildAgentSkills(rootAgent), + Capabilities: a2acore.AgentCapabilities{Streaming: true}, } - router.Handle(a2asrv.WellKnownAgentCardPath, a2asrv.NewStaticAgentCardHandler(agentCard)) + + compatProducer := a2av0.NewStaticAgentCardProducer(agentCard) + router.Handle(a2asrv.WellKnownAgentCardPath, a2asrv.NewAgentCardHandler(compatProducer)) agent := config.AgentLoader.RootAgent() executor := adka2a.NewExecutor(adka2a.ExecutorConfig{ RunnerConfig: runner.Config{ AppName: agent.Name(), Agent: agent, + MemoryService: config.MemoryService, SessionService: config.SessionService, ArtifactService: config.ArtifactService, PluginConfig: config.PluginConfig, }, }) reqHandler := a2asrv.NewHandler(executor, config.A2AOptions...) + router.Handle(apiPath, a2asrv.NewJSONRPCHandler(reqHandler)) + router.Handle(compatAPIPath, a2av0.NewJSONRPCHandler(reqHandler)) + return nil } diff --git a/cmd/launcher/web/a2a/a2a_test.go b/cmd/launcher/web/a2a/a2a_test.go index 147a495d1..77b436838 100644 --- a/cmd/launcher/web/a2a/a2a_test.go +++ b/cmd/launcher/web/a2a/a2a_test.go @@ -21,9 +21,12 @@ import ( "testing" "time" - a2acore "github.com/a2aproject/a2a-go/a2a" - "github.com/a2aproject/a2a-go/a2aclient" - "github.com/a2aproject/a2a-go/a2aclient/agentcard" + a2alegacy "github.com/a2aproject/a2a-go/a2a" + a2alegacyclient "github.com/a2aproject/a2a-go/a2aclient" + a2alegacyagentcard "github.com/a2aproject/a2a-go/a2aclient/agentcard" + a2acore "github.com/a2aproject/a2a-go/v2/a2a" + "github.com/a2aproject/a2a-go/v2/a2aclient" + "github.com/a2aproject/a2a-go/v2/a2aclient/agentcard" "google.golang.org/genai" "google.golang.org/adk/agent" @@ -72,7 +75,7 @@ func TestWebLauncher_ServesA2A(t *testing.T) { Name: "HelloWorldAgent", Run: func(ic agent.InvocationContext) iter.Seq2[*session.Event, error] { return func(yield func(*session.Event, error) bool) { - event := session.NewEvent(ic.InvocationID()) + event := session.NewEventWithContext(ic, ic.InvocationID()) event.Content = genai.NewContentFromText(wantMessage, genai.RoleModel) yield(event, nil) } @@ -92,41 +95,83 @@ func TestWebLauncher_ServesA2A(t *testing.T) { } }() - var card *a2acore.AgentCard - for retry := range 3 { - time.Sleep(10 * time.Millisecond) // give server time to start - card, err = agentcard.DefaultResolver.Resolve(ctx, "http://localhost:"+strconv.Itoa(port)) - if err == nil { - break + t.Run("A2A v2 client", func(t *testing.T) { + var card *a2acore.AgentCard + for retry := range 3 { + time.Sleep(10 * time.Millisecond) // give server time to start + card, err = agentcard.DefaultResolver.Resolve(ctx, "http://localhost:"+strconv.Itoa(port)) + if err == nil { + break + } + if retry == 2 { + t.Fatalf("cardResolver.Resolve() error = %v", err) + } } - if retry == 2 { - t.Fatalf("cardResolver.Resolve() error = %v", err) + + client, err := a2aclient.NewFromCard(ctx, card) + if err != nil { + t.Fatalf("a2aclient.NewFromCard() error = %v", err) } - } - client, err := a2aclient.NewFromCard(ctx, card) - if err != nil { - t.Fatalf("a2aclient.NewFromCard() error = %v", err) - } + got, err := client.SendMessage(ctx, &a2acore.SendMessageRequest{ + Message: a2acore.NewMessage(a2acore.MessageRoleUser, a2acore.NewTextPart("Hi!")), + }) + if err != nil { + t.Fatalf("client.SendMessage() error = %v", err) + } + task, ok := got.(*a2acore.Task) + if !ok { + t.Fatalf("client.SendMessage() result type = %T, want a2a.Task", got) + } + if len(task.Artifacts) != 1 { + t.Fatalf("len(task.Artifacts) = %d, want 1", len(task.Artifacts)) + } + parts := task.Artifacts[0].Parts + if len(parts) != 1 { + t.Fatalf("len(task.Artifacts[0].Parts) = %d, want 1", len(parts)) + } + if gotPart := parts[0].Text(); gotPart != wantMessage { + t.Fatalf("task.Artifacts[0].Parts[0] = %v, want %v", parts[0], wantMessage) + } + }) + + t.Run("A2A v0 client", func(t *testing.T) { + var card *a2alegacy.AgentCard + for retry := range 3 { + time.Sleep(10 * time.Millisecond) // give server time to start + card, err = a2alegacyagentcard.DefaultResolver.Resolve(ctx, "http://localhost:"+strconv.Itoa(port)) + if err == nil { + break + } + if retry == 2 { + t.Fatalf("a2alegacyagentcard.DefaultResolver.Resolve() error = %v", err) + } + } + + client, err := a2alegacyclient.NewFromCard(ctx, card) + if err != nil { + t.Fatalf("a2alegacyclient.NewFromCard() error = %v", err) + } - got, err := client.SendMessage(ctx, &a2acore.MessageSendParams{ - Message: a2acore.NewMessage(a2acore.MessageRoleUser, a2acore.TextPart{Text: "Hi!"}), + got, err := client.SendMessage(ctx, &a2alegacy.MessageSendParams{ + Message: a2alegacy.NewMessage(a2alegacy.MessageRoleUser, a2alegacy.TextPart{Text: "Hi!"}), + }) + if err != nil { + t.Fatalf("client.SendMessage() error = %v", err) + } + task, ok := got.(*a2alegacy.Task) + if !ok { + t.Fatalf("client.SendMessage() result type = %T, want a2alegacy.Task", got) + } + if len(task.Artifacts) != 1 { + t.Fatalf("len(task.Artifacts) = %d, want 1", len(task.Artifacts)) + } + parts := task.Artifacts[0].Parts + if len(parts) != 1 { + t.Fatalf("len(task.Artifacts[0].Parts) = %d, want 1", len(parts)) + } + if gotPart := parts[0].(a2alegacy.TextPart); gotPart.Text != wantMessage { + t.Fatalf("task.Artifacts[0].Parts[0] = %v, want %v", parts[0], wantMessage) + } }) - if err != nil { - t.Fatalf("client.SendMessage() error = %v", err) - } - task, ok := got.(*a2acore.Task) - if !ok { - t.Fatalf("client.SendMessage() result type = %T, want a2a.Task", got) - } - if len(task.Artifacts) != 1 { - t.Fatalf("len(task.Artifacts) = %d, want 1", len(task.Artifacts)) - } - parts := task.Artifacts[0].Parts - if len(parts) != 1 { - t.Fatalf("len(task.Artifacts[0].Parts) = %d, want 1", len(parts)) - } - if gotPart, ok := parts[0].(a2acore.TextPart); !ok || gotPart.Text != wantMessage { - t.Fatalf("task.Artifacts[0].Parts[0] = %v, want %v", parts[0], a2acore.TextPart{Text: wantMessage}) - } } diff --git a/examples/a2a/main.go b/examples/a2a/main.go index 824f3311a..7e74ccbda 100644 --- a/examples/a2a/main.go +++ b/examples/a2a/main.go @@ -23,18 +23,18 @@ import ( "net/url" "os" - "github.com/a2aproject/a2a-go/a2a" - "github.com/a2aproject/a2a-go/a2asrv" + "github.com/a2aproject/a2a-go/v2/a2a" + "github.com/a2aproject/a2a-go/v2/a2asrv" "google.golang.org/genai" "google.golang.org/adk/agent" "google.golang.org/adk/agent/llmagent" - "google.golang.org/adk/agent/remoteagent" + "google.golang.org/adk/agent/remoteagent/v2" "google.golang.org/adk/cmd/launcher" "google.golang.org/adk/cmd/launcher/full" "google.golang.org/adk/model/gemini" "google.golang.org/adk/runner" - "google.golang.org/adk/server/adka2a" + "google.golang.org/adk/server/adka2a/v2" "google.golang.org/adk/session" "google.golang.org/adk/tool" "google.golang.org/adk/tool/geminitool" @@ -42,7 +42,7 @@ import ( // newWeatherAgent creates a simple LLM-agent as in the quickstart example. func newWeatherAgent(ctx context.Context) agent.Agent { - model, err := gemini.NewModel(ctx, "gemini-2.5-flash", &genai.ClientConfig{ + model, err := gemini.NewModel(ctx, "gemini-3.1-flash-lite", &genai.ClientConfig{ APIKey: os.Getenv("GOOGLE_API_KEY"), }) if err != nil { @@ -79,10 +79,19 @@ func startWeatherAgentServer() string { agentPath := "/invoke" agentCard := &a2a.AgentCard{ - Name: agent.Name(), + Name: agent.Name(), + Description: agent.Description(), + SupportedInterfaces: []*a2a.AgentInterface{ + { + URL: baseURL.JoinPath(agentPath).String(), + ProtocolBinding: a2a.TransportProtocolJSONRPC, + ProtocolVersion: a2a.Version, + }, + }, + Version: "1.0.0", + DefaultInputModes: []string{"text/plain"}, + DefaultOutputModes: []string{"text/plain"}, Skills: adka2a.BuildAgentSkills(agent), - PreferredTransport: a2a.TransportProtocolJSONRPC, - URL: baseURL.JoinPath(agentPath).String(), Capabilities: a2a.AgentCapabilities{Streaming: true}, } @@ -113,8 +122,8 @@ func main() { a2aServerAddress := startWeatherAgentServer() remoteAgent, err := remoteagent.NewA2A(remoteagent.A2AConfig{ - Name: "A2A Weather agent", - AgentCardSource: a2aServerAddress, + Name: "A2A Weather agent", + AgentCardProvider: remoteagent.NewAgentCardProvider(a2aServerAddress), }) if err != nil { log.Fatalf("Failed to create a remote agent: %v", err) diff --git a/examples/agentengine/main.go b/examples/agentengine/main.go index 6c87c3b79..9f56ef889 100644 --- a/examples/agentengine/main.go +++ b/examples/agentengine/main.go @@ -17,6 +17,7 @@ package main import ( "context" + "fmt" "log" "math/rand/v2" "os" @@ -27,12 +28,51 @@ import ( "google.golang.org/adk/agent/llmagent" "google.golang.org/adk/cmd/launcher" "google.golang.org/adk/cmd/launcher/agentengine" + vertexaiMem "google.golang.org/adk/memory/vertexai" "google.golang.org/adk/model/gemini" + "google.golang.org/adk/plugin" + "google.golang.org/adk/runner" "google.golang.org/adk/session/vertexai" "google.golang.org/adk/tool" "google.golang.org/adk/tool/functiontool" + vertexaiutil "google.golang.org/adk/util/vertexai" ) +// Args defines the input structure for the memory search tool. +type Args struct { + Query string `json:"query" jsonschema:"The query to search for in the memory."` +} + +// Result defines the output structure for the memory search tool. +type Result struct { + Results []string `json:"results"` +} + +const ( + stateKeySessionLastUpdateTime = "sessionLastUpdateTime" +) + +// memorySearchToolFunc is the implementation of the memory search tool. +// This function demonstrates accessing memory via agent.ToolContext. +func memorySearchToolFunc(tctx agent.ToolContext, args Args) (Result, error) { + // The SearchMemory function is available on the context. + searchResults, err := tctx.SearchMemory(tctx, args.Query) + if err != nil { + log.Printf("Error searching memory: %v", err) + return Result{}, fmt.Errorf("failed memory search: %w", err) + } + + var results []string + for _, res := range searchResults.Memories { + if res.Content != nil { + for _, part := range res.Content.Parts { + results = append(results, part.Text) + } + } + } + return Result{Results: results}, nil +} + func main() { ctx := context.Background() @@ -58,7 +98,7 @@ func main() { type Output struct { Result int `json:"result"` } - handler := func(ctx tool.Context, input Input) (Output, error) { + handler := func(ctx agent.ToolContext, input Input) (Output, error) { return Output{ Result: input.Min + rand.IntN(input.Max-input.Min+1), }, nil @@ -71,13 +111,26 @@ func main() { log.Fatalf("Failed to create tool: %v", err) } + // Define a tool that can search memory. + memorySearchTool, err := functiontool.New( + functiontool.Config{ + Name: "search_past_conversations", + Description: "Searches past conversations for relevant information.", + }, + memorySearchToolFunc, + ) + if err != nil { + log.Fatalf("Failed to create tool: %v", err) + } + a, err := llmagent.New(llmagent.Config{ Name: "ae_agent", Model: model, Description: "General helpful agent", - Instruction: "You are a helpful agent, you should answer any questions you are given. Use 'random' tool to provide random numbers.", + Instruction: "You are a helpful agent, you should answer any questions you are given. Use 'random' tool to provide random numbers. Use search_past_conversations tool to get the facts about the user", Tools: []tool.Tool{ randomTool, + memorySearchTool, }, }) if err != nil { @@ -94,9 +147,55 @@ func main() { log.Fatalf("Failed to create session service: %v", err) } + memService, err := vertexaiMem.NewService(ctx, + &vertexaiMem.ServiceConfig{ + AgentEngineData: vertexaiutil.AgentEngineData{ + ProjectID: projectID, + Location: location, + ReasoningEngine: agentEngineID, + }, + StateKeySessionLastUpdateTime: stateKeySessionLastUpdateTime, + }) + if err != nil { + log.Fatalf("Failed to create memory service: %v", err) + } + + memPlugin, err := plugin.New(plugin.Config{ + Name: "Memory generator", + BeforeRunCallback: func(ic agent.InvocationContext) (*genai.Content, error) { + state := ic.Session().State() + err := state.Set(stateKeySessionLastUpdateTime, ic.Session().LastUpdateTime()) + if err != nil { + log.Printf("state.Set failed: %v\n", err) + return nil, err + } + return nil, nil + }, + AfterRunCallback: func(ic agent.InvocationContext) { + m := ic.Memory() + if m == nil { + log.Printf("ic.Memory() is nil\n") + return + } + err := m.AddSessionToMemory(ic, ic.Session()) + if err != nil { + log.Printf("ic.Memory().AddSessionToMemory failed: %v\n", err) + } + }, + }) + if err != nil { + log.Fatalf("Failed to create plugin: %v", err) + } + config := &launcher.Config{ SessionService: sessionService, AgentLoader: agent.NewSingleLoader(a), + MemoryService: memService, + PluginConfig: runner.PluginConfig{ + Plugins: []*plugin.Plugin{ + memPlugin, + }, + }, } l := agentengine.NewLauncher(agentEngineID) diff --git a/examples/bidi/README.md b/examples/bidi/README.md new file mode 100644 index 000000000..87d4b152f --- /dev/null +++ b/examples/bidi/README.md @@ -0,0 +1,114 @@ +# ADK Go Bidirectional Streaming Demo + +This directory contains a working demonstration of real-time bidirectional streaming using the Agent Development Kit (ADK) for Go. It showcases how to set up an agent with tools and serve it over a WebSocket connection for real-time interaction. + +![bidi-demo-screen](assets/bidi-demo-screen.png) + +## Overview + +The example in `main.go` sets up: +1. An **LLM Agent** named `bidi-demo` using the `gemini-3.1-flash-live-preview` model. +2. **Tools**: The agent is equipped with Google Search and a custom `camera_toggle` function tool. +3. An **HTTP Server**: It serves a web interface and handles WebSocket connections for the bidirectional streaming API. + +## Features + +- **Bidirectional Streaming**: Real-time communication handled by ADK's `RuntimeAPIController`. +- **Function Calling**: Demonstrates how to register and use custom Go functions as tools. +- **Web UI**: A simple frontend to interact with the agent, located in the `static/` directory. + +## Architecture + +The application uses ADK's `RuntimeAPIController` to manage the bidirectional streaming session: + +``` +┌─────────────┐ ┌──────────────────────┐ ┌─────────────┐ +│ │ │ │ │ │ +│ WebSocket │────────▶│ RuntimeAPIController │────────▶│ Live API │ +│ Client │ │ (RunLiveHandler) │ │ Session │ +│ │◀────────│ │◀────────│ │ +└─────────────┘ └──────────────────────┘ └─────────────┘ +``` + +- **Upstream**: The client sends audio, text, or media messages over the WebSocket connection. +- **Downstream**: The controller streams model responses and events back to the client in real-time. + +## Prerequisites + +- **Go**: Ensure you have Go installed (Go 1.23+ recommended). +- **API Key**: You need a Google API Key to access the Gemini models. + +## Getting Started + +### 1. Set up your API Key + +Set the `GOOGLE_API_KEY` environment variable: + +```bash +export GOOGLE_API_KEY="your_api_key_here" +``` + +### 2. Run the Server + +You can run the server from the project root or directly from this directory: + +**From the project root:** +```bash +go run examples/bidi/main.go +``` + +**From the `examples/bidi` directory:** +```bash +go run main.go +``` + +The server will start and serve the UI on `http://localhost:8081`. + +### 3. Access the UI + +Open your browser and navigate to `http://localhost:8081`. You should see a chat interface where you can interact with the agent. + +## Project Structure + +- `main.go`: The main entry point that configures a simple agent and starts the server. +- `static/`: Contains the frontend files (HTML, CSS, JS) shared by all examples. +- `streamingtool/`: Demonstrates a **streaming tool** that yields data over time (counting with delays) and how to handle the `stop_streaming` control signal. + - **Run**: `go run examples/bidi/streamingtool/main.go` +- `sequential/`: Demonstrates a **Sequential Agent** flow where control is passed from an 'Idea Generator' agent to a 'Story Teller' agent. + - **Run**: `go run examples/bidi/sequential/main.go` + +## Code Overview + +The core setup in `main.go` involves: + +- Creating the model: + ```go + model, err := gemini.NewModel(ctx, "gemini-3.1-flash-live-preview", &genai.ClientConfig{ + APIKey: os.Getenv("GOOGLE_API_KEY"), + }) + ``` +- Defining a custom tool: + ```go + cameraTool, err := functiontool.New(functiontool.Config{ + Name: "camera_toggle", + Description: "Turns the camera on or off.", + }, func(ctx tool.Context, args EmptyArgs) (MessageResult, error) { + // ... + }) + ``` +- Initializing the agent: + ```go + a, err := llmagent.New(llmagent.Config{ + Name: "bidi-demo", + Model: model, + Instruction: "You are a real-time voice assistant.", + Tools: []tool.Tool{geminitool.GoogleSearch{}, cameraTool}, + }) + ``` +- Serving the UI and the Live Handler: + ```go + controller := controllers.NewRuntimeAPIController(ss, nil, agent.NewSingleLoader(a), nil, 0, runner.PluginConfig{}, true) + http.HandleFunc("/run_live", func(w http.ResponseWriter, req *http.Request) { + controller.RunLiveHandler(w, req) + }) + ``` diff --git a/examples/bidi/assets/bidi-demo-screen.png b/examples/bidi/assets/bidi-demo-screen.png new file mode 100644 index 000000000..31e2b79f3 Binary files /dev/null and b/examples/bidi/assets/bidi-demo-screen.png differ diff --git a/examples/bidi/main.go b/examples/bidi/main.go new file mode 100644 index 000000000..68b6f8545 --- /dev/null +++ b/examples/bidi/main.go @@ -0,0 +1,104 @@ +// Copyright 2025 Google LLC +// +// Licensed under the Apache License, Version 2.0 (the "License"); +// you may not use this file except in compliance with the License. +// You may obtain a copy of the License at +// +// http://www.apache.org/licenses/LICENSE-2.0 +// +// Unless required by applicable law or agreed to in writing, software +// distributed under the License is distributed on an "AS IS" BASIS, +// WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. +// See the License for the specific language governing permissions and +// limitations under the License. + +// Package provides a quickstart ADK agent. +package main + +import ( + "context" + "fmt" + "log" + "net/http" + "os" + "path/filepath" + "runtime" + + "google.golang.org/genai" + + "google.golang.org/adk/agent" + "google.golang.org/adk/agent/llmagent" + "google.golang.org/adk/model/gemini" + "google.golang.org/adk/runner" + "google.golang.org/adk/server/adkrest/controllers" + "google.golang.org/adk/session" + "google.golang.org/adk/tool" + "google.golang.org/adk/tool/functiontool" + "google.golang.org/adk/tool/geminitool" +) + +func main() { + log.SetOutput(os.Stdout) + ctx := context.Background() + + model, err := gemini.NewModel(ctx, "gemini-3.1-flash-live-preview", &genai.ClientConfig{ + APIKey: os.Getenv("GOOGLE_API_KEY"), + }) + if err != nil { + log.Fatalf("Failed to create model: %v", err) + } + + type EmptyArgs struct{} + type MessageResult struct { + Message string `json:"message"` + } + + cameraTool, err := functiontool.New(functiontool.Config{ + Name: "camera_toggle", + Description: "Turns the camera on or off.", + }, func(ctx agent.ToolContext, args EmptyArgs) (MessageResult, error) { + fmt.Println("Camera tool was called!") + return MessageResult{Message: "Camera tool called successfully!"}, nil + }) + if err != nil { + log.Fatalf("Failed to create camera tool: %v", err) + } + + a, err := llmagent.New(llmagent.Config{ + Name: "bidi-demo", + Model: model, + Description: "Agent optimized for real-time bidirectional streaming.", + Instruction: "You are a real-time voice assistant.", + Tools: []tool.Tool{ + geminitool.GoogleSearch{}, + cameraTool, + }, + }) + if err != nil { + log.Fatalf("Failed to create agent: %v", err) + } + + // Create runner + ss := session.InMemoryService() + + _, filename, _, ok := runtime.Caller(0) + if !ok { + log.Fatal("No caller information") + } + staticDir := filepath.Join(filepath.Dir(filename), "static") + fs := http.FileServer(http.Dir(staticDir)) + http.Handle("/", fs) + http.Handle("/static/", http.StripPrefix("/static/", fs)) + + controller := controllers.NewRuntimeAPIController(ss, nil, agent.NewSingleLoader(a), nil, 0, runner.PluginConfig{}, true) + + http.HandleFunc("/run_live", func(w http.ResponseWriter, req *http.Request) { + err := controller.RunLiveHandler(w, req) + if err != nil { + log.Printf("RunLiveHandler failed: %v", err) + } + }) + + fmt.Println("Serving UI on http://localhost:8081") + log.Fatal(http.ListenAndServe(":8081", nil)) +} diff --git a/examples/bidi/sequential/main.go b/examples/bidi/sequential/main.go new file mode 100644 index 000000000..8631c68bb --- /dev/null +++ b/examples/bidi/sequential/main.go @@ -0,0 +1,111 @@ +// Copyright 2026 Google LLC +// +// Licensed under the Apache License, Version 2.0 (the "License"); +// you may not use this file except in compliance with the License. +// You may obtain a copy of the License at +// +// http://www.apache.org/licenses/LICENSE-2.0 +// +// Unless required by applicable law or agreed to in writing, software +// distributed under the License is distributed on an "AS IS" BASIS, +// WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. +// See the License for the specific language governing permissions and +// limitations under the License. + +// Package main provides an example of using sequential agents with real-time bidirectional streaming. +package main + +import ( + "context" + "fmt" + "log" + "net/http" + "os" + "path/filepath" + "runtime" + + "google.golang.org/genai" + + "google.golang.org/adk/agent" + "google.golang.org/adk/agent/llmagent" + "google.golang.org/adk/agent/workflowagents/sequentialagent" + "google.golang.org/adk/model/gemini" + "google.golang.org/adk/runner" + "google.golang.org/adk/server/adkrest/controllers" + "google.golang.org/adk/session" + "google.golang.org/adk/tool" + "google.golang.org/adk/tool/geminitool" +) + +func main() { + log.SetOutput(os.Stdout) + ctx := context.Background() + + model, err := gemini.NewModel(ctx, "gemini-2.5-flash-native-audio-preview-12-2025", &genai.ClientConfig{ + APIKey: os.Getenv("GOOGLE_API_KEY"), + }) + if err != nil { + log.Fatalf("Failed to create model: %v", err) + } + + ideaGenerator, err := llmagent.New(llmagent.Config{ + Name: "idea_generator", + Model: model, + Description: "Brainstorms creative story ideas with the user.", + Instruction: "You are the Idea Generator. Ask the user for a topic they are interested in, and brainstorm 3 creative story ideas for them. Discuss the options and help them choose their favorite idea. Once the user confirms their choice, call the task_completed function so the Story Teller agent can take over to narrate the story.", + Tools: []tool.Tool{ + geminitool.GoogleSearch{}, + }, + }) + if err != nil { + log.Fatalf("Failed to create idea generator agent: %v", err) + } + + storyTeller, err := llmagent.New(llmagent.Config{ + Name: "story_teller", + Model: model, + Description: "Narrates an engaging story based on the chosen idea.", + // Note: Sending content history is currently not implemented for the Live API. + // Therefore, the agent asks the user to remind them of the chosen idea instead of reviewing history. + Instruction: "You are the Story Teller. The previous agent has just finalized a story idea with the user. Greet the user, ask them to remind you of the chosen idea, and then narrate an exciting, highly engaging short story about it using your voice.", + Tools: []tool.Tool{ + geminitool.GoogleSearch{}, + }, + }) + if err != nil { + log.Fatalf("Failed to create story teller agent: %v", err) + } + + seqAgent, err := sequentialagent.New(sequentialagent.Config{ + AgentConfig: agent.Config{ + Name: "bidi-demo", + SubAgents: []agent.Agent{ideaGenerator, storyTeller}, + }, + }) + if err != nil { + log.Fatalf("Failed to create sequential agent: %v", err) + } + + ss := session.InMemoryService() + + _, filename, _, ok := runtime.Caller(0) + if !ok { + log.Fatal("No caller information") + } + staticDir := filepath.Join(filepath.Dir(filename), "../static") + fs := http.FileServer(http.Dir(staticDir)) + http.Handle("/", fs) + http.Handle("/static/", http.StripPrefix("/static/", fs)) + + controller := controllers.NewRuntimeAPIController(ss, nil, agent.NewSingleLoader(seqAgent), nil, 0, runner.PluginConfig{}, true) + + http.HandleFunc("/run_live", func(w http.ResponseWriter, req *http.Request) { + err := controller.RunLiveHandler(w, req) + if err != nil { + log.Printf("RunLiveHandler failed: %v", err) + } + }) + + fmt.Println("Serving UI on http://localhost:8081") + log.Fatal(http.ListenAndServe(":8081", nil)) +} diff --git a/examples/bidi/static/css/style.css b/examples/bidi/static/css/style.css new file mode 100644 index 000000000..4105b0739 --- /dev/null +++ b/examples/bidi/static/css/style.css @@ -0,0 +1,951 @@ +/** + * Modern Chat UI Styles for ADK Streaming Demo + */ + +:root { + --primary-color: #4285f4; + --user-bubble-bg: #4285f4; + --agent-bubble-bg: #f1f3f4; + --user-text-color: #ffffff; + --agent-text-color: #202124; + --bg-color: #ffffff; + --border-color: #e0e0e0; + --input-bg: #f8f9fa; + --shadow: 0 1px 3px rgba(0, 0, 0, 0.1); + --shadow-lg: 0 4px 12px rgba(0, 0, 0, 0.15); +} + +* { + margin: 0; + padding: 0; + box-sizing: border-box; +} + +body { + font-family: -apple-system, BlinkMacSystemFont, 'Segoe UI', Roboto, 'Helvetica Neue', Arial, sans-serif; + background-color: var(--bg-color); + color: #202124; + line-height: 1.6; + height: 100vh; + display: flex; + flex-direction: column; +} + +/* Header */ +header { + background: linear-gradient(135deg, #667eea 0%, #764ba2 100%); + color: white; + padding: 1.5rem 2rem; + box-shadow: var(--shadow-lg); + position: relative; +} + +h1 { + font-size: 1.5rem; + font-weight: 600; + margin: 0; +} + +.subtitle { + font-size: 0.875rem; + opacity: 0.9; + margin-top: 0.25rem; +} + +/* Header Options (Proactivity, Affective Dialog checkboxes) */ +.header-options { + display: flex; + align-items: center; + gap: 1.25rem; + margin-top: 0.75rem; +} + +.header-checkbox { + display: flex; + align-items: center; + gap: 0.375rem; + font-size: 0.8125rem; + cursor: pointer; + user-select: none; + opacity: 0.9; + transition: opacity 0.2s ease; +} + +.header-checkbox:hover { + opacity: 1; +} + +.header-checkbox input[type="checkbox"] { + width: 16px; + height: 16px; + cursor: pointer; + accent-color: #ffffff; +} + +.header-checkbox span { + white-space: nowrap; +} + +.connection-status { + position: absolute; + top: 1.5rem; + right: 2rem; + display: flex; + align-items: center; + gap: 0.5rem; + font-size: 0.875rem; +} + +.status-indicator { + width: 8px; + height: 8px; + border-radius: 50%; + background-color: #34a853; +} + +.status-indicator.disconnected { + background-color: #ea4335; +} + +@keyframes pulse { + 0%, 100% { + opacity: 1; + } + 50% { + opacity: 0.5; + } +} + +/* Main Layout: Split view for chat and console */ +.main-layout { + flex: 1; + display: flex; + max-width: 1800px; + width: 100%; + margin: 0 auto; + overflow: hidden; + gap: 0; +} + +/* Main Container: Chat area (2/3 of the layout) */ +.container { + flex: 2; + display: flex; + flex-direction: column; + overflow: hidden; + border-right: 1px solid var(--border-color); +} + +/* Messages Area */ +#messages { + flex: 1; + overflow-y: auto; + padding: 2rem; + display: flex; + flex-direction: column; + gap: 1rem; + background: linear-gradient(to bottom, #f8f9fa 0%, #ffffff 100%); +} + +/* Scroll styling */ +#messages::-webkit-scrollbar { + width: 8px; +} + +#messages::-webkit-scrollbar-track { + background: transparent; +} + +#messages::-webkit-scrollbar-thumb { + background: #dadce0; + border-radius: 4px; +} + +#messages::-webkit-scrollbar-thumb:hover { + background: #bdc1c6; +} + +/* Message Bubbles */ +.message { + display: flex; + margin-bottom: 0.5rem; + animation: slideIn 0.3s ease-out; +} + +@keyframes slideIn { + from { + opacity: 0; + transform: translateY(10px); + } + to { + opacity: 1; + transform: translateY(0); + } +} + +.message.user { + justify-content: flex-end; +} + +.message.agent { + justify-content: flex-start; +} + +.bubble { + max-width: 70%; + padding: 0.75rem 1rem; + border-radius: 1.25rem; + word-wrap: break-word; + box-shadow: var(--shadow); + position: relative; +} + +.message.user .bubble { + background-color: var(--user-bubble-bg); + color: var(--user-text-color); + border-bottom-right-radius: 0.25rem; +} + +.message.agent .bubble { + background-color: var(--agent-bubble-bg); + color: var(--agent-text-color); + border-bottom-left-radius: 0.25rem; +} + +.bubble-text { + margin: 0; + line-height: 1.5; +} + +/* Interrupted message styling */ +.message.interrupted .bubble { + opacity: 0.6; + background-color: #e8eaed; + border-left: 3px solid #f4b400; +} + +.message.interrupted .bubble::after { + content: 'interrupted'; + display: block; + font-size: 0.75rem; + color: #5f6368; + font-style: italic; + margin-top: 0.25rem; +} + +/* Transcription message styling */ +.message.transcription.user .bubble { + opacity: 0.9; + border: 1px solid rgba(255, 255, 255, 0.3); +} + +.message.transcription.user .bubble::before { + content: '🎤'; + opacity: 0.8; + margin-right: 0.25rem; +} + +/* Typing indicator */ +.typing-indicator { + display: inline-block; + margin-left: 0.25rem; + color: #5f6368; +} + +.typing-indicator::after { + content: '...'; + animation: ellipsis 1.5s infinite; +} + +@keyframes ellipsis { + 0%, 20% { + content: '.'; + } + 40% { + content: '..'; + } + 60%, 100% { + content: '...'; + } +} + +/* Image bubble styling */ +.bubble.image-bubble { + padding: 0.25rem; + max-width: 80%; +} + +.bubble-image { + max-width: 100%; + max-height: 300px; + width: auto; + height: auto; + border-radius: 0.75rem; + display: block; + object-fit: contain; +} + +/* Input Form */ +.input-container { + border-top: 1px solid var(--border-color); + background-color: white; + padding: 1.5rem 2rem; + box-shadow: 0 -2px 8px rgba(0, 0, 0, 0.05); +} + +#messageForm { + display: flex; + gap: 1rem; + align-items: center; +} + +.input-wrapper { + flex: 1; + display: flex; + gap: 0.75rem; + flex-wrap: wrap; +} + +#message { + flex: 1; + padding: 0.875rem 1rem; + border: 1px solid var(--border-color); + border-radius: 1.5rem; + font-size: 1rem; + font-family: inherit; + background-color: var(--input-bg); + transition: all 0.2s ease; + outline: none; +} + +#message:focus { + border-color: var(--primary-color); + background-color: white; + box-shadow: 0 0 0 3px rgba(66, 133, 244, 0.1); +} + +/* Buttons */ +button { + padding: 0.875rem 1.5rem; + border: none; + border-radius: 1.5rem; + font-size: 0.9375rem; + font-weight: 500; + cursor: pointer; + transition: all 0.2s ease; + font-family: inherit; + white-space: nowrap; +} + +#sendButton { + background-color: var(--primary-color); + color: white; +} + +#sendButton:hover:not(:disabled) { + background-color: #3367d6; + box-shadow: var(--shadow); + transform: translateY(-1px); +} + +#sendButton:disabled { + background-color: #dadce0; + color: #80868b; + cursor: not-allowed; + opacity: 0.6; +} + +#startAudioButton { + background-color: #34a853; + color: white; +} + +#startAudioButton:hover:not(:disabled) { + background-color: #2d8e47; + box-shadow: var(--shadow); + transform: translateY(-1px); +} + +#startAudioButton:disabled { + background-color: #dadce0; + color: #80868b; + cursor: not-allowed; + opacity: 0.6; +} + +#cameraButton { + background-color: #ea4335; + color: white; +} + +#cameraButton:hover:not(:disabled) { + background-color: #d33426; + box-shadow: var(--shadow); + transform: translateY(-1px); +} + +#cameraButton:disabled { + background-color: #dadce0; + color: #80868b; + cursor: not-allowed; + opacity: 0.6; +} + +#sendFileButton { + background-color: #a142f4; + color: white; +} + +#sendFileButton:hover:not(:disabled) { + background-color: #8b14f4; + box-shadow: var(--shadow); + transform: translateY(-1px); +} + +#sendFileButton:disabled { + background-color: #dadce0; + color: #80868b; + cursor: not-allowed; + opacity: 0.6; +} + +/* System Messages */ +.system-message { + text-align: center; + color: #5f6368; + font-size: 0.875rem; + padding: 0.5rem; + margin: 1rem 0; + font-style: italic; +} + +/* Console Panel (1/3 of the layout) */ +.console-panel { + flex: 1; + display: flex; + flex-direction: column; + background-color: #1e1e1e; + color: #d4d4d4; + font-family: 'Monaco', 'Menlo', 'Ubuntu Mono', 'Consolas', 'source-code-pro', monospace; + overflow: hidden; +} + +.console-header { + display: flex; + justify-content: space-between; + align-items: center; + padding: 0.75rem 1rem; + background-color: #2d2d2d; + border-bottom: 1px solid #3e3e3e; +} + +.console-header h2 { + font-size: 0.875rem; + font-weight: 600; + margin: 0; + color: #cccccc; + text-transform: uppercase; + letter-spacing: 0.5px; +} + +.console-controls { + display: flex; + align-items: center; + gap: 0.75rem; +} + +.console-checkbox { + display: flex; + align-items: center; + gap: 0.375rem; + font-size: 0.75rem; + color: #999999; + cursor: pointer; + user-select: none; +} + +.console-checkbox input[type="checkbox"] { + width: 14px; + height: 14px; + cursor: pointer; + accent-color: #4285f4; +} + +.console-checkbox span { + white-space: nowrap; +} + +.console-clear-btn { + padding: 0.375rem 0.75rem; + font-size: 0.75rem; + background-color: #3e3e3e; + color: #cccccc; + border: 1px solid #4e4e4e; + border-radius: 0.25rem; + cursor: pointer; + transition: all 0.2s ease; +} + +.console-clear-btn:hover { + background-color: #4e4e4e; + border-color: #5e5e5e; +} + +.console-content { + flex: 1; + overflow-y: auto; + padding: 0.75rem; + font-size: 0.75rem; + line-height: 1.5; +} + +/* Console entry */ +.console-entry { + margin-bottom: 0.75rem; + padding: 0.5rem; + border-left: 3px solid transparent; + background-color: rgba(255, 255, 255, 0.06); + border-radius: 0.25rem; + transition: background-color 0.2s ease; +} + +.console-entry.outgoing { + border-left-color: #4285f4; +} + +.console-entry.incoming { + border-left-color: #34a853; +} + +.console-entry.error { + border-left-color: #ea4335; + background-color: rgba(234, 67, 53, 0.15); +} + +/* Expandable console entries */ +.console-entry.expandable { + cursor: pointer; +} + +.console-entry.expandable:hover { + background-color: rgba(255, 255, 255, 0.10); +} + +.console-entry.expanded { + background-color: rgba(255, 255, 255, 0.08); +} + +.console-entry-header { + display: flex; + justify-content: space-between; + align-items: center; + margin-bottom: 0.375rem; +} + +.console-entry-left { + display: flex; + align-items: center; + gap: 0.5rem; +} + +.console-entry-emoji { + font-size: 0.9rem; + line-height: 1; + display: inline-block; + user-select: none; + min-width: 16px; + text-align: center; +} + +.console-expand-icon { + font-size: 0.6rem; + color: #858585; + width: 12px; + display: inline-block; + transition: transform 0.2s ease; + user-select: none; +} + +.console-entry-type { + font-weight: 600; + font-size: 0.7rem; + text-transform: uppercase; + letter-spacing: 0.5px; +} + +.console-entry.outgoing .console-entry-type { + color: #4285f4; +} + +.console-entry.incoming .console-entry-type { + color: #34a853; +} + +.console-entry.error .console-entry-type { + color: #ea4335; +} + +.console-entry-author { + font-size: 0.65rem; + font-weight: 500; + padding: 0.125rem 0.375rem; + border-radius: 0.25rem; + text-transform: lowercase; + letter-spacing: 0.3px; + border: 1px solid; +} + +/* Default style (for agents) */ +.console-entry-author { + background-color: rgba(156, 220, 254, 0.15); + color: #9cdcfe; + border-color: rgba(156, 220, 254, 0.3); +} + +/* User author badge */ +.console-entry-author[data-author="user"] { + background-color: rgba(66, 133, 244, 0.2); + color: #80b3ff; + border-color: rgba(66, 133, 244, 0.4); +} + +/* System author badge */ +.console-entry-author[data-author="system"] { + background-color: rgba(133, 133, 133, 0.2); + color: #b0b0b0; + border-color: rgba(133, 133, 133, 0.3); +} + +.console-entry-timestamp { + color: #858585; + font-size: 0.65rem; +} + +.console-entry-content { + color: #d4d4d4; + white-space: pre-wrap; + word-break: break-word; + font-size: 0.7rem; + line-height: 1.4; + padding-left: 2.5rem; /* Indent to align with header after emoji and icon */ +} + +.console-entry-content:empty { + display: none; +} + +/* Style for quoted text in summaries (transcriptions, text responses) */ +.console-entry-content::first-line { + color: #e0e0e0; +} + +.console-entry-json { + background-color: #252526; + padding: 0.5rem; + border-radius: 0.25rem; + margin-top: 0.5rem; + overflow-x: auto; + max-height: 400px; + overflow-y: auto; + transition: all 0.3s ease; +} + +.console-entry-json.collapsed { + display: none; +} + +.console-entry-json pre { + margin: 0; + color: #9cdcfe; +} + +/* Highlight key fields in JSON */ +.json-key { + color: #9cdcfe; +} + +.json-string { + color: #ce9178; +} + +.json-number { + color: #b5cea8; +} + +.json-boolean { + color: #569cd6; +} + +.json-null { + color: #569cd6; +} + +/* Console scrollbar */ +.console-content::-webkit-scrollbar { + width: 8px; +} + +.console-content::-webkit-scrollbar-track { + background: #1e1e1e; +} + +.console-content::-webkit-scrollbar-thumb { + background: #3e3e3e; + border-radius: 4px; +} + +.console-content::-webkit-scrollbar-thumb:hover { + background: #4e4e4e; +} + +/* JSON scrollbar */ +.console-entry-json::-webkit-scrollbar { + width: 6px; + height: 6px; +} + +.console-entry-json::-webkit-scrollbar-track { + background: #1e1e1e; +} + +.console-entry-json::-webkit-scrollbar-thumb { + background: #3e3e3e; + border-radius: 3px; +} + +.console-entry-json::-webkit-scrollbar-thumb:hover { + background: #4e4e4e; +} + +/* Responsive Design */ +@media (max-width: 768px) { + header { + padding: 1rem 1.5rem; + } + + h1 { + font-size: 1.25rem; + } + + .connection-status { + position: static; + margin-top: 0.5rem; + } + + /* Stack console panel below chat on mobile */ + .main-layout { + flex-direction: column; + } + + .console-panel { + max-height: 300px; + border-top: 1px solid var(--border-color); + } + + .container { + border-right: none; + } + + #messages { + padding: 1rem; + } + + .bubble { + max-width: 85%; + } + + .input-container { + padding: 1rem; + } + + #messageForm { + flex-direction: column; + gap: 0.75rem; + } + + .input-wrapper { + width: 100%; + flex-direction: column; + } + + button { + width: 100%; + } +} + +/* Loading state */ +.loading { + display: inline-block; + width: 20px; + height: 20px; + border: 3px solid rgba(255, 255, 255, 0.3); + border-radius: 50%; + border-top-color: white; + animation: spin 0.8s linear infinite; +} + +@keyframes spin { + to { + transform: rotate(360deg); + } +} + +/* Video bubble styling */ +.bubble.video-bubble { + padding: 0.25rem; + max-width: 80%; + background-color: #000; +} + +.bubble-video { + max-width: 100%; + max-height: 300px; + width: auto; + height: auto; + border-radius: 0.75rem; + display: block; + object-fit: contain; +} + +/* Camera Preview Modal */ +.modal { + display: none; + position: fixed; + z-index: 1000; + left: 0; + top: 0; + width: 100%; + height: 100%; + background-color: rgba(0, 0, 0, 0.7); + backdrop-filter: blur(4px); +} + +.modal.show { + display: flex; + align-items: center; + justify-content: center; +} + +.modal-content { + background-color: white; + border-radius: 1rem; + box-shadow: 0 8px 32px rgba(0, 0, 0, 0.3); + max-width: 90%; + max-height: 90%; + width: 640px; + display: flex; + flex-direction: column; + animation: modalSlideIn 0.3s ease-out; +} + +@keyframes modalSlideIn { + from { + opacity: 0; + transform: translateY(-20px); + } + to { + opacity: 1; + transform: translateY(0); + } +} + +.modal-header { + display: flex; + justify-content: space-between; + align-items: center; + padding: 1.5rem; + border-bottom: 1px solid var(--border-color); +} + +.modal-header h3 { + margin: 0; + font-size: 1.25rem; + font-weight: 600; + color: #202124; +} + +.close-btn { + background: none; + border: none; + font-size: 2rem; + line-height: 1; + color: #5f6368; + cursor: pointer; + padding: 0; + width: 32px; + height: 32px; + display: flex; + align-items: center; + justify-content: center; + border-radius: 50%; + transition: all 0.2s ease; +} + +.close-btn:hover { + background-color: #f1f3f4; + color: #202124; +} + +.modal-body { + padding: 1.5rem; + flex: 1; + display: flex; + align-items: center; + justify-content: center; + background-color: #000; + border-radius: 0 0 1rem 1rem; +} + +#cameraPreview { + width: 100%; + max-height: 480px; + border-radius: 0.5rem; + object-fit: contain; +} + +.modal-footer { + padding: 1.5rem; + border-top: 1px solid var(--border-color); + display: flex; + gap: 1rem; + justify-content: flex-end; + background-color: white; + border-radius: 0 0 1rem 1rem; +} + +.btn-primary { + background-color: var(--primary-color); + color: white; + padding: 0.875rem 1.5rem; + border: none; + border-radius: 1.5rem; + font-size: 0.9375rem; + font-weight: 500; + cursor: pointer; + transition: all 0.2s ease; + font-family: inherit; +} + +.btn-primary:hover { + background-color: #3367d6; + box-shadow: var(--shadow); + transform: translateY(-1px); +} + +.btn-secondary { + background-color: white; + color: #5f6368; + padding: 0.875rem 1.5rem; + border: 1px solid var(--border-color); + border-radius: 1.5rem; + font-size: 0.9375rem; + font-weight: 500; + cursor: pointer; + transition: all 0.2s ease; + font-family: inherit; +} + +.btn-secondary:hover { + background-color: #f8f9fa; + border-color: #bdc1c6; +} \ No newline at end of file diff --git a/examples/bidi/static/index.html b/examples/bidi/static/index.html new file mode 100644 index 000000000..5ac581073 --- /dev/null +++ b/examples/bidi/static/index.html @@ -0,0 +1,89 @@ + + + + + + ADK Gemini Live API Toolkit Demo + + + + + +
+

ADK Gemini Live API Toolkit Demo

+
Real-time bidirectional streaming with Google ADK
+
+ + +
+
+ + Connecting... +
+
+ +
+
+
+ + +
+
+
+ + + + + + + +
+
+
+
+ +
+
+

Event Console

+
+ + +
+
+
+
+
+ + + + + \ No newline at end of file diff --git a/examples/bidi/static/js/app.js b/examples/bidi/static/js/app.js new file mode 100644 index 000000000..52b09716c --- /dev/null +++ b/examples/bidi/static/js/app.js @@ -0,0 +1,1304 @@ +/** + * app.js: JS code for the ADK Gemini Live API Toolkit demo app. + */ + +/** + * WebSocket handling + */ + +// Connect the server with a WebSocket connection +const userId = "demo-user"; +//let sessionId = "demo-session-" + Math.random().toString(36).substring(7); +let sessionId = "fixed-demo-session"; +let websocket = null; +let is_audio = false; +let audioBuffer = []; + +// Get checkbox elements for RunConfig options +const enableProactivityCheckbox = document.getElementById("enableProactivity"); +const enableAffectiveDialogCheckbox = document.getElementById("enableAffectiveDialog"); + +// Reconnect WebSocket when RunConfig options change +function handleRunConfigChange() { + if (websocket && websocket.readyState === WebSocket.OPEN) { + addSystemMessage("Reconnecting with updated settings..."); + addConsoleEntry('outgoing', 'Reconnecting due to settings change', { + proactivity: enableProactivityCheckbox.checked, + affective_dialog: enableAffectiveDialogCheckbox.checked + }, '🔄', 'system'); + + // Keep same session ID for testing history + // sessionId = "demo-session-" + Math.random().toString(36).substring(7); + + websocket.close(); + // connectWebsocket() will be called by onclose handler after delay + } +} + +// Add change listeners to RunConfig checkboxes +enableProactivityCheckbox.addEventListener("change", handleRunConfigChange); +enableAffectiveDialogCheckbox.addEventListener("change", handleRunConfigChange); + +// Build WebSocket URL with RunConfig options as query parameters +function getWebSocketUrl() { + // Use wss:// for HTTPS pages, ws:// for HTTP (localhost development) + const wsProtocol = window.location.protocol === "https:" ? "wss:" : "ws:"; + const baseUrl = wsProtocol + "//" + window.location.host + "/run_live"; + const params = new URLSearchParams(); + + // Add required parameters for RunLiveHandler + params.append("appName", "bidi-demo"); + params.append("userId", userId); + params.append("sessionId", sessionId); + + // Add proactivity option if checked + if (enableProactivityCheckbox && enableProactivityCheckbox.checked) { + params.append("proactivity", "true"); + } + + // Add affective dialog option if checked + if (enableAffectiveDialogCheckbox && enableAffectiveDialogCheckbox.checked) { + params.append("affective_dialog", "true"); + } + + const queryString = params.toString(); + return queryString ? baseUrl + "?" + queryString : baseUrl; +} + +// Get DOM elements +const messageForm = document.getElementById("messageForm"); +const messageInput = document.getElementById("message"); +const messagesDiv = document.getElementById("messages"); +const statusIndicator = document.getElementById("statusIndicator"); +const statusText = document.getElementById("statusText"); +const consoleContent = document.getElementById("consoleContent"); +const clearConsoleBtn = document.getElementById("clearConsole"); +const showAudioEventsCheckbox = document.getElementById("showAudioEvents"); +let currentMessageId = null; +let currentBubbleElement = null; +let currentInputTranscriptionId = null; +let currentInputTranscriptionElement = null; +let currentOutputTranscriptionId = null; +let currentOutputTranscriptionElement = null; +let lastAgentBubbleElement = null; +let inputTranscriptionFinished = false; // Track if input transcription is complete for this turn +let hasOutputTranscriptionInTurn = false; // Track if output transcription delivered the response + +// Helper function to clean spaces between CJK characters +// Removes spaces between Japanese/Chinese/Korean characters while preserving spaces around Latin text +function cleanCJKSpaces(text) { + // CJK Unicode ranges: Hiragana, Katakana, Kanji, CJK Unified Ideographs, Fullwidth forms + const cjkPattern = /[\u3000-\u303f\u3040-\u309f\u30a0-\u30ff\u4e00-\u9faf\uff00-\uffef]/; + + // Remove spaces between two CJK characters + return text.replace(/(\S)\s+(?=\S)/g, (match, char1) => { + // Get the character after the space(s) + const nextCharMatch = text.match(new RegExp(char1 + '\\s+(.)', 'g')); + if (nextCharMatch && nextCharMatch.length > 0) { + const char2 = nextCharMatch[0].slice(-1); + // If both characters are CJK, remove the space + if (cjkPattern.test(char1) && cjkPattern.test(char2)) { + return char1; + } + } + return match; + }); +} + +// Console logging functionality +function formatTimestamp() { + const now = new Date(); + return now.toLocaleTimeString('en-US', { hour12: false, hour: '2-digit', minute: '2-digit', second: '2-digit', fractionalSecondDigits: 3 }); +} + +function addConsoleEntry(type, content, data = null, emoji = null, author = null, isAudio = false) { + // Skip audio events if checkbox is unchecked + if (isAudio && !showAudioEventsCheckbox.checked) { + return; + } + + const entry = document.createElement("div"); + entry.className = `console-entry ${type}`; + + const header = document.createElement("div"); + header.className = "console-entry-header"; + + const leftSection = document.createElement("div"); + leftSection.className = "console-entry-left"; + + // Add emoji icon if provided + if (emoji) { + const emojiIcon = document.createElement("span"); + emojiIcon.className = "console-entry-emoji"; + emojiIcon.textContent = emoji; + leftSection.appendChild(emojiIcon); + } + + // Add expand/collapse icon + const expandIcon = document.createElement("span"); + expandIcon.className = "console-expand-icon"; + expandIcon.textContent = data ? "▶" : ""; + + const typeLabel = document.createElement("span"); + typeLabel.className = "console-entry-type"; + typeLabel.textContent = type === 'outgoing' ? '↑ Upstream' : type === 'incoming' ? '↓ Downstream' : '⚠ Error'; + + leftSection.appendChild(expandIcon); + leftSection.appendChild(typeLabel); + + // Add author badge if provided + if (author) { + const authorBadge = document.createElement("span"); + authorBadge.className = "console-entry-author"; + authorBadge.textContent = author; + authorBadge.setAttribute('data-author', author); + leftSection.appendChild(authorBadge); + } + + const timestamp = document.createElement("span"); + timestamp.className = "console-entry-timestamp"; + timestamp.textContent = formatTimestamp(); + + header.appendChild(leftSection); + header.appendChild(timestamp); + + const contentDiv = document.createElement("div"); + contentDiv.className = "console-entry-content"; + contentDiv.textContent = content; + + entry.appendChild(header); + entry.appendChild(contentDiv); + + // JSON details (hidden by default) + let jsonDiv = null; + if (data) { + jsonDiv = document.createElement("div"); + jsonDiv.className = "console-entry-json collapsed"; + const pre = document.createElement("pre"); + pre.textContent = JSON.stringify(data, null, 2); + jsonDiv.appendChild(pre); + entry.appendChild(jsonDiv); + + // Make entry clickable if it has data + entry.classList.add("expandable"); + + // Toggle expand/collapse on click + entry.addEventListener("click", () => { + const isExpanded = !jsonDiv.classList.contains("collapsed"); + + if (isExpanded) { + // Collapse + jsonDiv.classList.add("collapsed"); + expandIcon.textContent = "▶"; + entry.classList.remove("expanded"); + } else { + // Expand + jsonDiv.classList.remove("collapsed"); + expandIcon.textContent = "▼"; + entry.classList.add("expanded"); + } + }); + } + + consoleContent.appendChild(entry); + consoleContent.scrollTop = consoleContent.scrollHeight; +} + +function clearConsole() { + consoleContent.innerHTML = ''; +} + +// Clear console button handler +clearConsoleBtn.addEventListener('click', clearConsole); + +// Update connection status UI +function updateConnectionStatus(connected) { + if (connected) { + statusIndicator.classList.remove("disconnected"); + statusText.textContent = "Connected"; + } else { + statusIndicator.classList.add("disconnected"); + statusText.textContent = "Disconnected"; + } +} + +// Create a message bubble element +function createMessageBubble(text, isUser, isPartial = false) { + const messageDiv = document.createElement("div"); + messageDiv.className = `message ${isUser ? "user" : "agent"}`; + + const bubbleDiv = document.createElement("div"); + bubbleDiv.className = "bubble"; + + const textP = document.createElement("p"); + textP.className = "bubble-text"; + textP.textContent = text; + + // Add typing indicator for partial messages + if (isPartial && !isUser) { + const typingSpan = document.createElement("span"); + typingSpan.className = "typing-indicator"; + textP.appendChild(typingSpan); + } + + bubbleDiv.appendChild(textP); + messageDiv.appendChild(bubbleDiv); + + return messageDiv; +} + +// Create an image message bubble element +function createImageBubble(imageDataUrl, isUser) { + const messageDiv = document.createElement("div"); + messageDiv.className = `message ${isUser ? "user" : "agent"}`; + + const bubbleDiv = document.createElement("div"); + bubbleDiv.className = "bubble image-bubble"; + + const img = document.createElement("img"); + img.src = imageDataUrl; + img.className = "bubble-image"; + img.alt = "Captured image"; + + bubbleDiv.appendChild(img); + messageDiv.appendChild(bubbleDiv); + + return messageDiv; +} + +// Update existing message bubble text +function updateMessageBubble(element, text, isPartial = false) { + const textElement = element.querySelector(".bubble-text"); + + // Remove existing typing indicator + const existingIndicator = textElement.querySelector(".typing-indicator"); + if (existingIndicator) { + existingIndicator.remove(); + } + + textElement.textContent = text; + + // Add typing indicator for partial messages + if (isPartial) { + const typingSpan = document.createElement("span"); + typingSpan.className = "typing-indicator"; + textElement.appendChild(typingSpan); + } +} + +// Add a system message +function addSystemMessage(text) { + const messageDiv = document.createElement("div"); + messageDiv.className = "system-message"; + messageDiv.textContent = text; + appendMessage(messageDiv); + scrollToBottom(); +} + +// Scroll to bottom of messages +function scrollToBottom() { + messagesDiv.scrollTop = messagesDiv.scrollHeight; +} + +// Append message to messagesDiv, inserting before stream bubble if active +function appendMessage(element) { + const streamBubble = document.getElementById("streamPreviewBubble"); + if (streamBubble) { + messagesDiv.insertBefore(element, streamBubble); + } else { + messagesDiv.appendChild(element); + } +} + +// Sanitize event data for console display (replace large audio data with summary) +function sanitizeEventForDisplay(event) { + // Deep clone the event object + const sanitized = JSON.parse(JSON.stringify(event)); + + // Check for audio data in content.parts + if (sanitized.content && sanitized.content.parts) { + sanitized.content.parts = sanitized.content.parts.map(part => { + if (part.inlineData && part.inlineData.data) { + // Calculate byte size (base64 string length / 4 * 3, roughly) + const byteSize = Math.floor(part.inlineData.data.length * 0.75); + return { + ...part, + inlineData: { + ...part.inlineData, + data: `(${byteSize.toLocaleString()} bytes)` + } + }; + } + return part; + }); + } + + return sanitized; +} + +// WebSocket handlers +function connectWebsocket() { + // Connect websocket + const ws_url = getWebSocketUrl(); + websocket = new WebSocket(ws_url); + + // Handle connection open + websocket.onopen = function () { + console.log("WebSocket connection opened."); + updateConnectionStatus(true); + addSystemMessage("Connected to ADK streaming server"); + + // Log to console + addConsoleEntry('incoming', 'WebSocket Connected', { + userId: userId, + sessionId: sessionId, + url: ws_url + }, '🔌', 'system'); + + // Enable the Send button + document.getElementById("sendButton").disabled = false; + addSubmitHandler(); + }; + + // Handle incoming messages + websocket.onmessage = function (event) { + // Parse the incoming ADK Event + const adkEvent = JSON.parse(event.data); + console.log("[AGENT TO CLIENT] ", adkEvent); + + // Log to console panel + let eventSummary = 'Event'; + let eventEmoji = '📨'; // Default emoji + const author = adkEvent.author || 'system'; + + if (adkEvent.turnComplete) { + eventSummary = 'Turn Complete'; + eventEmoji = '✅'; + } else if (adkEvent.interrupted) { + eventSummary = 'Interrupted'; + eventEmoji = '⏸️'; + } else if (adkEvent.inputTranscription) { + // Show transcription text in summary + const transcriptionText = adkEvent.inputTranscription.text || ''; + const truncated = transcriptionText.length > 60 + ? transcriptionText.substring(0, 60) + '...' + : transcriptionText; + eventSummary = `Input Transcription: "${truncated}"`; + eventEmoji = '📝'; + } else if (adkEvent.outputTranscription) { + // Show transcription text in summary + const transcriptionText = adkEvent.outputTranscription.text || ''; + const truncated = transcriptionText.length > 60 + ? transcriptionText.substring(0, 60) + '...' + : transcriptionText; + eventSummary = `Output Transcription: "${truncated}"`; + eventEmoji = '📝'; + } else if (adkEvent.usageMetadata) { + // Show token usage information + const usage = adkEvent.usageMetadata; + const promptTokens = usage.promptTokenCount || 0; + const responseTokens = usage.candidatesTokenCount || 0; + const totalTokens = usage.totalTokenCount || 0; + eventSummary = `Token Usage: ${totalTokens.toLocaleString()} total (${promptTokens.toLocaleString()} prompt + ${responseTokens.toLocaleString()} response)`; + eventEmoji = '📊'; + } else if (adkEvent.content && adkEvent.content.parts) { + const hasText = adkEvent.content.parts.some(p => p.text); + const hasAudio = adkEvent.content.parts.some(p => p.inlineData); + const hasExecutableCode = adkEvent.content.parts.some(p => p.executableCode); + const hasCodeExecutionResult = adkEvent.content.parts.some(p => p.codeExecutionResult); + + if (hasExecutableCode) { + // Show executable code + const codePart = adkEvent.content.parts.find(p => p.executableCode); + if (codePart && codePart.executableCode) { + const code = codePart.executableCode.code || ''; + const language = codePart.executableCode.language || 'unknown'; + const truncated = code.length > 60 + ? code.substring(0, 60).replace(/\n/g, ' ') + '...' + : code.replace(/\n/g, ' '); + eventSummary = `Executable Code (${language}): ${truncated}`; + eventEmoji = '💻'; + } + } + + if (hasCodeExecutionResult) { + // Show code execution result + const resultPart = adkEvent.content.parts.find(p => p.codeExecutionResult); + if (resultPart && resultPart.codeExecutionResult) { + const outcome = resultPart.codeExecutionResult.outcome || 'UNKNOWN'; + const output = resultPart.codeExecutionResult.output || ''; + const truncatedOutput = output.length > 60 + ? output.substring(0, 60).replace(/\n/g, ' ') + '...' + : output.replace(/\n/g, ' '); + eventSummary = `Code Execution Result (${outcome}): ${truncatedOutput}`; + eventEmoji = outcome === 'OUTCOME_OK' ? '✅' : '❌'; + } + } + + if (hasText) { + // Show text preview in summary + const textPart = adkEvent.content.parts.find(p => p.text); + if (textPart && textPart.text) { + const text = textPart.text; + const truncated = text.length > 80 + ? text.substring(0, 80) + '...' + : text; + eventSummary = `Text: "${truncated}"`; + eventEmoji = '💭'; + } else { + eventSummary = 'Text Response'; + eventEmoji = '💭'; + } + } + + if (hasAudio) { + // Extract audio info for summary + const audioPart = adkEvent.content.parts.find(p => p.inlineData); + if (audioPart && audioPart.inlineData) { + const mimeType = audioPart.inlineData.mimeType || 'unknown'; + const dataLength = audioPart.inlineData.data ? audioPart.inlineData.data.length : 0; + // Base64 string length / 4 * 3 gives approximate bytes + const byteSize = Math.floor(dataLength * 0.75); + eventSummary = `Audio Response: ${mimeType} (${byteSize.toLocaleString()} bytes)`; + eventEmoji = '🔊'; + } else { + eventSummary = 'Audio Response'; + eventEmoji = '🔊'; + } + + // Log audio event with isAudio flag (filtered by checkbox) + const sanitizedEvent = sanitizeEventForDisplay(adkEvent); + addConsoleEntry('incoming', eventSummary, sanitizedEvent, eventEmoji, author, true); + } + } + + // Create a sanitized version for console display (replace large audio data with summary) + // Skip if already logged as audio event above + const isAudioOnlyEvent = adkEvent.content && adkEvent.content.parts && + adkEvent.content.parts.some(p => p.inlineData) && + !adkEvent.content.parts.some(p => p.text); + if (!isAudioOnlyEvent) { + const sanitizedEvent = sanitizeEventForDisplay(adkEvent); + addConsoleEntry('incoming', eventSummary, sanitizedEvent, eventEmoji, author); + } + + // Handle turn complete event + if (adkEvent.turnComplete === true) { + // Remove typing indicator from current message + if (currentBubbleElement) { + const textElement = currentBubbleElement.querySelector(".bubble-text"); + const typingIndicator = textElement.querySelector(".typing-indicator"); + if (typingIndicator) { + typingIndicator.remove(); + } + } + // Remove typing indicator from current output transcription + if (currentOutputTranscriptionElement) { + const textElement = currentOutputTranscriptionElement.querySelector(".bubble-text"); + const typingIndicator = textElement.querySelector(".typing-indicator"); + if (typingIndicator) { + typingIndicator.remove(); + } + } + currentMessageId = null; + currentBubbleElement = null; + currentOutputTranscriptionId = null; + currentOutputTranscriptionElement = null; + inputTranscriptionFinished = false; // Reset for next turn + hasOutputTranscriptionInTurn = false; // Reset for next turn + return; + } + + // Handle interrupted event + if (adkEvent.interrupted === true) { + // Stop audio playback if it's playing + if (audioPlayerNode) { + audioPlayerNode.port.postMessage({ command: "endOfAudio" }); + } + + // Keep the partial message but mark it as interrupted + let markedInterrupted = false; + if (currentBubbleElement) { + const textElement = currentBubbleElement.querySelector(".bubble-text"); + + // Remove typing indicator + const typingIndicator = textElement.querySelector(".typing-indicator"); + if (typingIndicator) { + typingIndicator.remove(); + } + + // Add interrupted marker + currentBubbleElement.classList.add("interrupted"); + markedInterrupted = true; + } + + // Keep the partial output transcription but mark it as interrupted + if (currentOutputTranscriptionElement) { + const textElement = currentOutputTranscriptionElement.querySelector(".bubble-text"); + + // Remove typing indicator + const typingIndicator = textElement.querySelector(".typing-indicator"); + if (typingIndicator) { + typingIndicator.remove(); + } + + // Add interrupted marker + currentOutputTranscriptionElement.classList.add("interrupted"); + markedInterrupted = true; + } + + // Fallback to the last agent bubble element if the active trackers were already finalized + if (!markedInterrupted && lastAgentBubbleElement) { + lastAgentBubbleElement.classList.add("interrupted"); + } + + // Reset state so new content creates a new bubble + currentMessageId = null; + currentBubbleElement = null; + currentOutputTranscriptionId = null; + currentOutputTranscriptionElement = null; + inputTranscriptionFinished = false; // Reset for next turn + hasOutputTranscriptionInTurn = false; // Reset for next turn + return; + } + + // Handle input transcription (user's spoken words) + if (adkEvent.inputTranscription && adkEvent.inputTranscription.text) { + const transcriptionText = adkEvent.inputTranscription.text; + const isFinished = adkEvent.inputTranscription.finished; + + if (transcriptionText) { + // Ignore late-arriving transcriptions after we've finished for this turn + if (inputTranscriptionFinished) { + return; + } + + if (currentInputTranscriptionId == null) { + // Create new transcription bubble + currentInputTranscriptionId = Math.random().toString(36).substring(7); + // Clean spaces between CJK characters + const cleanedText = cleanCJKSpaces(transcriptionText); + currentInputTranscriptionElement = createMessageBubble(cleanedText, true, !isFinished); + currentInputTranscriptionElement.id = currentInputTranscriptionId; + + // Add a special class to indicate it's a transcription + currentInputTranscriptionElement.classList.add("transcription"); + + appendMessage(currentInputTranscriptionElement); + lastAgentBubbleElement = null; + } else { + // Update existing transcription bubble only if model hasn't started responding + // This prevents late partial transcriptions from overwriting complete ones + if (currentOutputTranscriptionId == null && currentMessageId == null) { + if (isFinished) { + // Final transcription contains the complete text, replace entirely + const cleanedText = cleanCJKSpaces(transcriptionText); + updateMessageBubble(currentInputTranscriptionElement, cleanedText, false); + } else { + // Partial transcription - append to existing text + if (currentInputTranscriptionElement) { + const existingText = currentInputTranscriptionElement.querySelector(".bubble-text").textContent; + // Remove typing indicator if present + const cleanText = existingText.replace(/\.\.\.$/, ''); + // Clean spaces between CJK characters before updating + const accumulatedText = cleanCJKSpaces(cleanText + transcriptionText); + updateMessageBubble(currentInputTranscriptionElement, accumulatedText, true); + } else { + console.log("fixed: currentInputTranscriptionElement was null in transcription handler"); + // Fallback: create a new bubble if it's null for some reason + currentInputTranscriptionId = Math.random().toString(36).substring(7); + const cleanedText = cleanCJKSpaces(transcriptionText); + currentInputTranscriptionElement = createMessageBubble(cleanedText, true, true); + currentInputTranscriptionElement.id = currentInputTranscriptionId; + currentInputTranscriptionElement.classList.add("transcription"); + appendMessage(currentInputTranscriptionElement); + } + } + } + } + + // If transcription is finished, reset the state and mark as complete + if (isFinished) { + currentInputTranscriptionId = null; + currentInputTranscriptionElement = null; + inputTranscriptionFinished = true; // Prevent duplicate bubbles from late events + } + + scrollToBottom(); + } + } + + // Handle output transcription (model's spoken words) + if (adkEvent.outputTranscription && adkEvent.outputTranscription.text) { + const transcriptionText = adkEvent.outputTranscription.text; + const isFinished = adkEvent.outputTranscription.finished; + hasOutputTranscriptionInTurn = true; + + if (transcriptionText) { + // Finalize any active input transcription when server starts responding + if (currentInputTranscriptionId != null && currentOutputTranscriptionId == null) { + // This is the first output transcription - finalize input transcription + if (currentInputTranscriptionElement) { + const textElement = currentInputTranscriptionElement.querySelector(".bubble-text"); + const typingIndicator = textElement.querySelector(".typing-indicator"); + if (typingIndicator) { + typingIndicator.remove(); + } + } else { + console.log("fixed: currentInputTranscriptionElement was null in content handler"); + } + // Reset input transcription state so next user input creates new balloon + currentInputTranscriptionId = null; + currentInputTranscriptionElement = null; + inputTranscriptionFinished = true; // Prevent duplicate bubbles from late events + } + + if (currentOutputTranscriptionId == null) { + // Create new transcription bubble for agent + currentOutputTranscriptionId = Math.random().toString(36).substring(7); + currentOutputTranscriptionElement = createMessageBubble(transcriptionText, false, !isFinished); + currentOutputTranscriptionElement.id = currentOutputTranscriptionId; + + // Add a special class to indicate it's a transcription + currentOutputTranscriptionElement.classList.add("transcription"); + + appendMessage(currentOutputTranscriptionElement); + lastAgentBubbleElement = currentOutputTranscriptionElement; + } else { + // Update existing transcription bubble + if (isFinished) { + // Final transcription contains the complete text, replace entirely + updateMessageBubble(currentOutputTranscriptionElement, transcriptionText, false); + } else { + // Partial transcription - append to existing text + const existingText = currentOutputTranscriptionElement.querySelector(".bubble-text").textContent; + // Remove typing indicator if present + const cleanText = existingText.replace(/\.\.\.$/, ''); + updateMessageBubble(currentOutputTranscriptionElement, cleanText + transcriptionText, true); + } + } + + // If transcription is finished, reset the state + if (isFinished) { + currentOutputTranscriptionId = null; + currentOutputTranscriptionElement = null; + } + + scrollToBottom(); + } + } + + // Handle content events (text or audio) + if (adkEvent.content && adkEvent.content.parts) { + const parts = adkEvent.content.parts; + + // Finalize any active input transcription when server starts responding with content + if (currentInputTranscriptionId != null && currentMessageId == null && currentOutputTranscriptionId == null) { + // This is the first content event - finalize input transcription + if (currentInputTranscriptionElement) { + const textElement = currentInputTranscriptionElement.querySelector(".bubble-text"); + const typingIndicator = textElement.querySelector(".typing-indicator"); + if (typingIndicator) { + typingIndicator.remove(); + } + } + // Reset input transcription state so next user input creates new balloon + currentInputTranscriptionId = null; + currentInputTranscriptionElement = null; + inputTranscriptionFinished = true; // Prevent duplicate bubbles from late events + } + + for (const part of parts) { + // Handle function response + if (part.functionResponse) { + console.log("Received function response:", part.functionResponse); + if (part.functionResponse.name === "camera_toggle") { + toggleVideoStreaming(); + } + } + + // Handle inline data (audio) + if (part.inlineData) { + const mimeType = part.inlineData.mimeType; + const data = part.inlineData.data; + + if (mimeType && mimeType.startsWith("audio/pcm")) { + const arrayBuffer = base64ToArray(data); + if (audioPlayerNode) { + audioPlayerNode.port.postMessage(arrayBuffer); + } else { + audioBuffer.push(arrayBuffer); + } + } + } + + // Handle text + if (part.text) { + // Skip thinking/reasoning text from chat bubbles (shown in event console) + if (part.thought) { + continue; + } + + // Skip final aggregated content when output transcription already + // delivered the response (prevents duplicate thinking text replay) + if (!adkEvent.partial && hasOutputTranscriptionInTurn) { + continue; + } + + // Handle system messages separately to avoid hijacking chat bubbles + if (adkEvent.author === "system") { + addSystemMessage(part.text); + continue; + } + + const isUser = (adkEvent.content.role === "user" || adkEvent.author === "user"); + + // Add a new message bubble for a new turn, or if role changed + if (currentMessageId == null || currentBubbleElement == null || currentBubbleElement.classList.contains("user") !== isUser) { + // Finalize previous bubble if role changed + if (currentBubbleElement && currentBubbleElement.classList.contains("user") !== isUser) { + const textElement = currentBubbleElement.querySelector(".bubble-text"); + const typingIndicator = textElement.querySelector(".typing-indicator"); + if (typingIndicator) { + typingIndicator.remove(); + } + } + + currentMessageId = Math.random().toString(36).substring(7); + currentBubbleElement = createMessageBubble(part.text, isUser, true); + currentBubbleElement.id = currentMessageId; + appendMessage(currentBubbleElement); + if (!isUser) { + lastAgentBubbleElement = currentBubbleElement; + } + } else { + // Update the existing message bubble with accumulated text + const existingText = currentBubbleElement.querySelector(".bubble-text").textContent; + // Remove the "..." if present + const cleanText = existingText.replace(/\.\.\.$/, ''); + updateMessageBubble(currentBubbleElement, cleanText + part.text, true); + } + + // Scroll down to the bottom of the messagesDiv + scrollToBottom(); + } + } + } + }; + + // Handle connection close + websocket.onclose = function (e) { + console.log("WebSocket connection closed.", e); + updateConnectionStatus(false); + document.getElementById("sendButton").disabled = true; + const reason = e && e.reason ? e.reason + " - " : ""; + addSystemMessage(`${reason}Connection closed. Reconnecting in 5 seconds...`); + + // Log to console + addConsoleEntry('error', 'WebSocket Disconnected', { + status: e && e.reason ? e.reason : "Connection closed", + code: e ? e.code : null, + reconnecting: true, + reconnectDelay: '5 seconds' + }, '🔌', 'system'); + + setTimeout(function () { + console.log("Reconnecting..."); + + // Log reconnection attempt to console + addConsoleEntry('outgoing', 'Reconnecting to ADK server...', { + userId: userId, + sessionId: sessionId + }, '🔄', 'system'); + + connectWebsocket(); + }, 5000); + }; + + websocket.onerror = function (e) { + console.log("WebSocket error: ", e); + updateConnectionStatus(false); + + // Log to console + addConsoleEntry('error', 'WebSocket Error', { + error: e.type, + message: 'Connection error occurred' + }, '⚠️', 'system'); + }; +} +connectWebsocket(); + +// Add submit handler to the form +function addSubmitHandler() { + messageForm.onsubmit = function (e) { + e.preventDefault(); + const message = messageInput.value.trim(); + if (message) { + // Add user message bubble + const userBubble = createMessageBubble(message, true, false); + appendMessage(userBubble); + scrollToBottom(); + + // Clear input + messageInput.value = ""; + + // Send message to server + sendMessage(message); + console.log("[CLIENT TO AGENT] " + message); + } + return false; + }; +} + +// Send a message to the server as JSON +function sendMessage(message) { + ensureAudioPlayerStarted(); + if (websocket && websocket.readyState == WebSocket.OPEN) { + lastAgentBubbleElement = null; + const jsonMessage = JSON.stringify({ + content: { + role: "user", + parts: [{ text: message }] + } + }); + websocket.send(jsonMessage); + + // Log to console panel + addConsoleEntry('outgoing', 'User Message: ' + message, null, '💬', 'user'); + } +} + +// Decode Base64 data to Array +// Handles both standard base64 and base64url encoding +function base64ToArray(base64) { + // Convert base64url to standard base64 + // Replace URL-safe characters: - with +, _ with / + let standardBase64 = base64.replace(/-/g, '+').replace(/_/g, '/'); + + // Add padding if needed + while (standardBase64.length % 4) { + standardBase64 += '='; + } + + const binaryString = window.atob(standardBase64); + const len = binaryString.length; + const bytes = new Uint8Array(len); + for (let i = 0; i < len; i++) { + bytes[i] = binaryString.charCodeAt(i); + } + return bytes.buffer; +} + +/** + * Camera handling + */ + +const cameraButton = document.getElementById("cameraButton"); +const streamVideoButton = document.getElementById("streamVideoButton"); +const cameraModal = document.getElementById("cameraModal"); +const cameraPreview = document.getElementById("cameraPreview"); +const closeCameraModal = document.getElementById("closeCameraModal"); +const cancelCamera = document.getElementById("cancelCamera"); +const captureImageBtn = document.getElementById("captureImage"); +const sendFileButton = document.getElementById("sendFileButton"); +const fileInput = document.getElementById("fileInput"); + +let cameraStream = null; +let isVideoStreaming = false; +let videoStreamInterval = null; + +// Open camera modal and start preview +async function openCameraPreview() { + try { + // Request access to the user's webcam + cameraStream = await navigator.mediaDevices.getUserMedia({ + video: { + width: { ideal: 768 }, + height: { ideal: 768 }, + facingMode: 'user' + } + }); + + // Set the stream to the video element + cameraPreview.srcObject = cameraStream; + + // Show the modal + cameraModal.classList.add('show'); + + } catch (error) { + console.error('Error accessing camera:', error); + addSystemMessage(`Failed to access camera: ${error.message}`); + + // Log to console + addConsoleEntry('error', 'Camera access failed', { + error: error.message, + name: error.name + }, '⚠️', 'system'); + } +} + +// Start video stream for preview in element +async function startVideoStream(videoElement) { + try { + cameraStream = await navigator.mediaDevices.getUserMedia({ + video: { + width: { ideal: 768 }, + height: { ideal: 768 }, + facingMode: 'user' + } + }); + + videoElement.srcObject = cameraStream; + + } catch (error) { + console.error('Error accessing camera for stream:', error); + addSystemMessage(`Failed to access camera: ${error.message}`); + addConsoleEntry('error', 'Camera access failed', { + error: error.message, + name: error.name + }, '⚠️', 'system'); + } +} + +// Close camera modal and stop preview +function closeCameraPreview() { + // Stop the camera stream + if (cameraStream) { + cameraStream.getTracks().forEach(track => track.stop()); + cameraStream = null; + } + + // Stop streaming if active + if (isVideoStreaming) { + clearInterval(videoStreamInterval); + isVideoStreaming = false; + streamVideoButton.textContent = "📹 Stream Video"; + streamVideoButton.classList.remove("active"); + addSystemMessage("Video streaming stopped"); + } + + // Clear the video source + cameraPreview.srcObject = null; + + // Hide the modal + cameraModal.classList.remove('show'); +} + +// Capture image from the live preview +function captureImageFromPreview() { + if (!cameraStream) { + addSystemMessage('No camera stream available'); + return; + } + + try { + // Create canvas to capture the frame + const canvas = document.createElement('canvas'); + canvas.width = cameraPreview.videoWidth; + canvas.height = cameraPreview.videoHeight; + const context = canvas.getContext('2d'); + + // Draw current video frame to canvas + context.drawImage(cameraPreview, 0, 0, canvas.width, canvas.height); + + // Convert canvas to data URL for display + const imageDataUrl = canvas.toDataURL('image/jpeg', 0.85); + + // Display the captured image in the chat + const imageBubble = createImageBubble(imageDataUrl, true); + appendMessage(imageBubble); + scrollToBottom(); + + // Convert canvas to blob for sending to server + canvas.toBlob((blob) => { + // Convert blob to base64 for sending to server + const reader = new FileReader(); + reader.onloadend = () => { + const base64data = reader.result.split(',')[1]; // Remove data:image/jpeg;base64, prefix + sendImage(base64data); + }; + reader.readAsDataURL(blob); + + // Log to console + addConsoleEntry('outgoing', `Image captured: ${blob.size} bytes (JPEG)`, { + size: blob.size, + type: 'image/jpeg', + dimensions: `${canvas.width}x${canvas.height}` + }, '📷', 'user'); + }, 'image/jpeg', 0.85); + + // Close the camera modal + closeCameraPreview(); + + } catch (error) { + console.error('Error capturing image:', error); + addSystemMessage(`Failed to capture image: ${error.message}`); + + // Log to console + addConsoleEntry('error', 'Image capture failed', { + error: error.message, + name: error.name + }, '⚠️', 'system'); + } +} + +// Send image to server +function sendImage(base64Image, mimeType = "image/jpeg") { + if (websocket && websocket.readyState === WebSocket.OPEN) { + const jsonMessage = JSON.stringify({ + blob: { + mime_type: mimeType, + data: base64Image + } + }); + websocket.send(jsonMessage); + console.log(`[CLIENT TO AGENT] Sent image (${mimeType})`); + } +} + +// Capture and send a single frame +function captureAndSendFrame() { + if (!cameraStream) return; + + try { + const canvas = document.createElement('canvas'); + // Use video in stream bubble if active, fallback to cameraPreview + const streamBubble = document.getElementById("streamPreviewBubble"); + const videoElement = streamBubble ? streamBubble.querySelector("video") : cameraPreview; + + canvas.width = videoElement.videoWidth; + canvas.height = videoElement.videoHeight; + const context = canvas.getContext('2d'); + context.drawImage(videoElement, 0, 0, canvas.width, canvas.height); + + canvas.toBlob((blob) => { + const reader = new FileReader(); + reader.onloadend = () => { + const base64data = reader.result.split(',')[1]; + sendImage(base64data); + }; + reader.readAsDataURL(blob); + }, 'image/jpeg', 0.6); // Lower quality for stream to save bandwidth + } catch (error) { + console.error('Error capturing video frame:', error); + } +} + +// Toggle video streaming +function toggleVideoStreaming() { + if (isVideoStreaming) { + // Stop streaming + clearInterval(videoStreamInterval); + isVideoStreaming = false; + streamVideoButton.textContent = "📹 Stream Video"; + streamVideoButton.classList.remove("active"); + addSystemMessage("Video streaming stopped"); + + // Keep the bubble in chat but stop tracks + const streamBubble = document.getElementById("streamPreviewBubble"); + if (streamBubble) { + // Remove ID so new messages don't get inserted before it anymore + streamBubble.removeAttribute("id"); + } + + if (cameraStream) { + cameraStream.getTracks().forEach(track => track.stop()); + cameraStream = null; + } + } else { + // Start streaming + // Create video bubble + const messageDiv = document.createElement("div"); + messageDiv.className = "message user"; + messageDiv.id = "streamPreviewBubble"; + + const bubbleDiv = document.createElement("div"); + bubbleDiv.className = "bubble video-bubble"; + + const video = document.createElement("video"); + video.className = "bubble-video"; + video.autoplay = true; + video.playsInline = true; + + bubbleDiv.appendChild(video); + messageDiv.appendChild(bubbleDiv); + + // Append to chat (always at bottom for now) + messagesDiv.appendChild(messageDiv); + scrollToBottom(); + + if (!cameraStream) { + startVideoStream(video).then(() => { + if (cameraStream) { + startStreamingLoop(); + } + }); + } else { + // Reuse existing stream if available + video.srcObject = cameraStream; + startStreamingLoop(); + } + } +} + +function startStreamingLoop() { + isVideoStreaming = true; + streamVideoButton.textContent = "⏹️ Stop Video"; + streamVideoButton.classList.add("active"); + addSystemMessage("Video streaming started"); + + // 1 FPS = 1000ms interval + videoStreamInterval = setInterval(captureAndSendFrame, 1000); +} + +// Event listeners +cameraButton.addEventListener("click", openCameraPreview); +streamVideoButton.addEventListener("click", toggleVideoStreaming); +closeCameraModal.addEventListener("click", closeCameraPreview); +cancelCamera.addEventListener("click", closeCameraPreview); +captureImageBtn.addEventListener("click", captureImageFromPreview); + +sendFileButton.addEventListener("click", () => { + fileInput.click(); +}); + +fileInput.addEventListener("change", (event) => { + const file = event.target.files[0]; + if (!file) return; + + const reader = new FileReader(); + reader.onloadend = () => { + const base64data = reader.result.split(',')[1]; + const mimeType = file.type; + + // Display the image in the chat + const imageBubble = createImageBubble(reader.result, true); + appendMessage(imageBubble); + scrollToBottom(); + + // Send to server + sendImage(base64data, mimeType); + + // Reset file input so the same file can be selected again + fileInput.value = ''; + + // Log to console + addConsoleEntry('outgoing', `File sent: ${file.name} (${file.size} bytes)`, { + name: file.name, + size: file.size, + type: mimeType + }, '📁', 'user'); + }; + reader.readAsDataURL(file); +}); + +// Close modal when clicking outside of it +cameraModal.addEventListener("click", (event) => { + if (event.target === cameraModal) { + closeCameraPreview(); + } +}); + +/** + * Audio handling + */ + +let audioPlayerNode; +let audioPlayerContext; +let audioRecorderNode; +let audioRecorderContext; +let micStream; + +// Import the audio worklets +import { startAudioPlayerWorklet } from "./audio-player.js"; +import { startAudioRecorderWorklet } from "./audio-recorder.js"; + +// Ensure audio player is started +function ensureAudioPlayerStarted() { + if (!audioPlayerContext) { + startAudioPlayerWorklet().then(([node, ctx]) => { + audioPlayerNode = node; + audioPlayerContext = ctx; + // Drain audio buffer + while (audioBuffer.length > 0) { + audioPlayerNode.port.postMessage(audioBuffer.shift()); + } + }); + } else if (audioPlayerContext.state === 'suspended') { + audioPlayerContext.resume(); + } +} + +// Start audio +function startAudio() { + // Start audio output + ensureAudioPlayerStarted(); + // Start audio input + return startAudioRecorderWorklet(audioRecorderHandler).then( + ([node, ctx, stream]) => { + audioRecorderNode = node; + audioRecorderContext = ctx; + micStream = stream; + } + ); +} + +// Toggle audio streaming +async function toggleAudioStreaming() { + if (is_audio) { + // Stop audio + is_audio = false; + if (micStream) { + micStream.getTracks().forEach(track => track.stop()); + micStream = null; + } + if (audioPlayerNode) { + audioPlayerNode.port.postMessage({ command: "endOfAudio" }); + } + if (audioRecorderContext) { + audioRecorderContext.close(); + audioRecorderContext = null; + } + // Keep audioPlayerContext open so we can still hear the agent + startAudioButton.textContent = "🎤 Start Voice"; + startAudioButton.classList.remove("active"); + addSystemMessage("Audio streaming stopped"); + addConsoleEntry('outgoing', 'Audio Mode Disabled', { status: 'Audio stopped' }, '🎤', 'system'); + } else { + // Start audio + try { + await startAudio(); + is_audio = true; + startAudioButton.textContent = "⏹️ Stop Voice"; + startAudioButton.classList.add("active"); + addSystemMessage("Audio mode enabled - you can now speak to the agent"); + addConsoleEntry('outgoing', 'Audio Mode Enabled', { + status: 'Audio worklets started', + message: 'Microphone active - audio input will be sent to agent' + }, '🎤', 'system'); + } catch (error) { + console.error("Failed to start audio:", error); + addSystemMessage(`Failed to start audio: ${error.message}`); + // Reset state + is_audio = false; + startAudioButton.textContent = "🎤 Start Voice"; + startAudioButton.classList.remove("active"); + addConsoleEntry('error', 'Failed to enable audio mode', { error: error.message }, '⚠️', 'system'); + } + } +} + +const startAudioButton = document.getElementById("startAudioButton"); +startAudioButton.addEventListener("click", toggleAudioStreaming); + +// Audio recorder handler +function audioRecorderHandler(pcmData) { + if (websocket && websocket.readyState === WebSocket.OPEN && is_audio) { + // Send audio as binary WebSocket frame (more efficient than base64 JSON) + websocket.send(pcmData); + console.log("[CLIENT TO AGENT] Sent audio chunk: %s bytes", pcmData.byteLength); + + // Log to console panel (optional, can be noisy with frequent audio chunks) + // addConsoleEntry('outgoing', `Audio chunk: ${pcmData.byteLength} bytes`); + } +} diff --git a/examples/bidi/static/js/audio-player.js b/examples/bidi/static/js/audio-player.js new file mode 100644 index 000000000..f7f235e01 --- /dev/null +++ b/examples/bidi/static/js/audio-player.js @@ -0,0 +1,24 @@ +/** + * Audio Player Worklet + */ + +export async function startAudioPlayerWorklet() { + // 1. Create an AudioContext + const audioContext = new AudioContext({ + sampleRate: 24000 + }); + + + // 2. Load your custom processor code + const workletURL = new URL('./pcm-player-processor.js', import.meta.url); + await audioContext.audioWorklet.addModule(workletURL); + + // 3. Create an AudioWorkletNode + const audioPlayerNode = new AudioWorkletNode(audioContext, 'pcm-player-processor'); + + // 4. Connect to the destination + audioPlayerNode.connect(audioContext.destination); + + // The audioPlayerNode.port is how we send messages (audio data) to the processor + return [audioPlayerNode, audioContext]; +} \ No newline at end of file diff --git a/examples/bidi/static/js/audio-recorder.js b/examples/bidi/static/js/audio-recorder.js new file mode 100644 index 000000000..876932a2f --- /dev/null +++ b/examples/bidi/static/js/audio-recorder.js @@ -0,0 +1,58 @@ +/** + * Audio Recorder Worklet + */ + +let micStream; + +export async function startAudioRecorderWorklet(audioRecorderHandler) { + // Create an AudioContext + const audioRecorderContext = new AudioContext({ sampleRate: 16000 }); + console.log("AudioContext sample rate:", audioRecorderContext.sampleRate); + + // Load the AudioWorklet module + const workletURL = new URL("./pcm-recorder-processor.js", import.meta.url); + await audioRecorderContext.audioWorklet.addModule(workletURL); + + // Request access to the microphone + micStream = await navigator.mediaDevices.getUserMedia({ + audio: { channelCount: 1 }, + }); + const source = audioRecorderContext.createMediaStreamSource(micStream); + + // Create an AudioWorkletNode that uses the PCMProcessor + const audioRecorderNode = new AudioWorkletNode( + audioRecorderContext, + "pcm-recorder-processor" + ); + + // Connect the microphone source to the worklet. + source.connect(audioRecorderNode); + audioRecorderNode.port.onmessage = (event) => { + // Convert to 16-bit PCM + const pcmData = convertFloat32ToPCM(event.data); + + // Send the PCM data to the handler. + audioRecorderHandler(pcmData); + }; + return [audioRecorderNode, audioRecorderContext, micStream]; +} + +/** + * Stop the microphone. + */ +export function stopMicrophone(micStream) { + micStream.getTracks().forEach((track) => track.stop()); + console.log("stopMicrophone(): Microphone stopped."); +} + +// Convert Float32 samples to 16-bit PCM. +function convertFloat32ToPCM(inputData) { + // Create an Int16Array of the same length. + const pcm16 = new Int16Array(inputData.length); + for (let i = 0; i < inputData.length; i++) { + // Multiply by 0x7fff (32767) to scale the float value to 16-bit PCM range. + pcm16[i] = inputData[i] * 0x7fff; + } + // Return the underlying ArrayBuffer. + return pcm16.buffer; +} \ No newline at end of file diff --git a/examples/bidi/static/js/pcm-player-processor.js b/examples/bidi/static/js/pcm-player-processor.js new file mode 100644 index 000000000..7e162362d --- /dev/null +++ b/examples/bidi/static/js/pcm-player-processor.js @@ -0,0 +1,75 @@ +/** + * An audio worklet processor that stores the PCM audio data sent from the main thread + * to a buffer and plays it. + */ +class PCMPlayerProcessor extends AudioWorkletProcessor { + constructor() { + super(); + + // Init buffer + this.bufferSize = 24000 * 180; // 24kHz x 180 seconds + this.buffer = new Float32Array(this.bufferSize); + this.writeIndex = 0; + this.readIndex = 0; + + // Handle incoming messages from main thread + this.port.onmessage = (event) => { + // Reset the buffer when 'endOfAudio' message received + if (event.data.command === 'endOfAudio') { + this.readIndex = this.writeIndex; // Clear the buffer + console.log("endOfAudio received, clearing the buffer."); + return; + } + + // Decode the base64 data to int16 array. + const int16Samples = new Int16Array(event.data); + + // Add the audio data to the buffer + this._enqueue(int16Samples); + }; + } + + // Push incoming Int16 data into our ring buffer. + _enqueue(int16Samples) { + for (let i = 0; i < int16Samples.length; i++) { + // Convert 16-bit integer to float in [-1, 1] + const floatVal = int16Samples[i] / 32768; + + // Store in ring buffer for left channel only (mono) + this.buffer[this.writeIndex] = floatVal; + this.writeIndex = (this.writeIndex + 1) % this.bufferSize; + + // Overflow handling (overwrite oldest samples) + if (this.writeIndex === this.readIndex) { + this.readIndex = (this.readIndex + 1) % this.bufferSize; + } + } + } + + // The system calls `process()` ~128 samples at a time (depending on the browser). + // We fill the output buffers from our ring buffer. + process(inputs, outputs, parameters) { + + // Write a frame to the output + const output = outputs[0]; + const framesPerBlock = output[0].length; + for (let frame = 0; frame < framesPerBlock; frame++) { + + // Write the sample(s) into the output buffer + output[0][frame] = this.buffer[this.readIndex]; // left channel + if (output.length > 1) { + output[1][frame] = this.buffer[this.readIndex]; // right channel + } + + // Move the read index forward unless underflowing + if (this.readIndex != this.writeIndex) { + this.readIndex = (this.readIndex + 1) % this.bufferSize; + } + } + + // Returning true tells the system to keep the processor alive + return true; + } +} + +registerProcessor('pcm-player-processor', PCMPlayerProcessor); diff --git a/examples/bidi/static/js/pcm-recorder-processor.js b/examples/bidi/static/js/pcm-recorder-processor.js new file mode 100644 index 000000000..585d67491 --- /dev/null +++ b/examples/bidi/static/js/pcm-recorder-processor.js @@ -0,0 +1,18 @@ +class PCMProcessor extends AudioWorkletProcessor { + constructor() { + super(); + } + + process(inputs, outputs, parameters) { + if (inputs.length > 0 && inputs[0].length > 0) { + // Use the first channel + const inputChannel = inputs[0][0]; + // Copy the buffer to avoid issues with recycled memory + const inputCopy = new Float32Array(inputChannel); + this.port.postMessage(inputCopy); + } + return true; + } +} + +registerProcessor("pcm-recorder-processor", PCMProcessor); \ No newline at end of file diff --git a/examples/bidi/streamingtool/main.go b/examples/bidi/streamingtool/main.go new file mode 100644 index 000000000..93d029f18 --- /dev/null +++ b/examples/bidi/streamingtool/main.go @@ -0,0 +1,143 @@ +// Copyright 2025 Google LLC +// +// Licensed under the Apache License, Version 2.0 (the "License"); +// you may not use this file except in compliance with the License. +// You may obtain a copy of the License at +// +// http://www.apache.org/licenses/LICENSE-2.0 +// +// Unless required by applicable law or agreed to in writing, software +// distributed under the License is distributed on an "AS IS" BASIS, +// WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. +// See the License for the specific language governing permissions and +// limitations under the License. + +// Package provides a quickstart ADK agent. +package main + +import ( + "context" + "fmt" + "iter" + "log" + "net/http" + "os" + "path/filepath" + "runtime" + "time" + + "google.golang.org/genai" + + "google.golang.org/adk/agent" + "google.golang.org/adk/agent/llmagent" + "google.golang.org/adk/model/gemini" + "google.golang.org/adk/runner" + "google.golang.org/adk/server/adkrest/controllers" + "google.golang.org/adk/session" + "google.golang.org/adk/tool" + "google.golang.org/adk/tool/functiontool" +) + +func main() { + log.SetOutput(os.Stdout) + ctx := context.Background() + + model, err := gemini.NewModel(ctx, "gemini-3.1-flash-live-preview", &genai.ClientConfig{ + APIKey: os.Getenv("GOOGLE_API_KEY"), + }) + if err != nil { + log.Fatalf("Failed to create model: %v", err) + } + + // Define a streaming tool that yields numbers. + counterTool, err := functiontool.NewStreaming(functiontool.Config{ + Name: "count_to", + Description: "Counts to a specified number, yielding each number with a delay.", + }, func(ctx agent.ToolContext, args struct { + N int `json:"n"` + }, + ) iter.Seq2[string, error] { + return func(yield func(string, error) bool) { + for i := 1; i <= args.N; i++ { + time.Sleep(5000 * time.Millisecond) + if !yield(fmt.Sprintf("Count: %d", i), nil) { + return + } + } + } + }) + if err != nil { + log.Fatalf("Failed to create streaming tool: %v", err) + } + + // Define a standard tool to stop streaming. + // Note: While defined here so that the model is aware of the tool declaration, + // during live bidirectional streaming the ADK Live Control Plane intercepts + // calls to "stop_streaming" dynamically. It will bulk-cancel the context for + // all running background goroutines executing that specific streaming tool name. + stopTool, err := functiontool.New(functiontool.Config{ + Name: "stop_streaming", + Description: "Stops a running streaming function.", + }, func(ctx agent.ToolContext, args struct { + FunctionName string `json:"function_name"` + }, + ) (map[string]any, error) { + return map[string]any{"status": fmt.Sprintf("Requested to stop %s", args.FunctionName)}, nil + }) + if err != nil { + log.Fatalf("Failed to create stop tool: %v", err) + } + + // Define a function tool to check divisibility. + checkDivisibleTool, err := functiontool.New(functiontool.Config{ + Name: "check_divisible", + Description: "Checks if a number is divisible by another number.", + }, func(ctx agent.ToolContext, args struct { + Number int `json:"number"` + Divisor int `json:"divisor"` + }, + ) (map[string]any, error) { + if args.Divisor == 0 { + return map[string]any{"result": false, "error": "cannot divide by zero"}, nil + } + fmt.Printf("Dividing %d by %d\n", args.Number, args.Divisor) + return map[string]any{"result": args.Number%args.Divisor == 0}, nil + }) + if err != nil { + log.Fatalf("Failed to create check divisible tool: %v", err) + } + + a, err := llmagent.New(llmagent.Config{ + Name: "bidi-demo", + Model: model, + Instruction: "You are a helpful assistant with a streaming tool 'count_to'. Always use it when asked to count. Wait for the tool results, and when you recieve them you should say the number, if it is divisible by 3 you should not say the number and instead say Fizz and if it is divisible by 5 you should say Buzz, if it is divisible by both 3 and 5 you should say FizzBuzz. Always use the check_divisible tool", + Tools: []tool.Tool{counterTool, stopTool, checkDivisibleTool}, + }) + if err != nil { + log.Fatalf("Failed to create agent: %v", err) + } + + // Create runner + ss := session.InMemoryService() + + _, filename, _, ok := runtime.Caller(0) + if !ok { + log.Fatal("No caller information") + } + staticDir := filepath.Join(filepath.Dir(filename), "../static") + fs := http.FileServer(http.Dir(staticDir)) + http.Handle("/", fs) + http.Handle("/static/", http.StripPrefix("/static/", fs)) + + controller := controllers.NewRuntimeAPIController(ss, nil, agent.NewSingleLoader(a), nil, 0, runner.PluginConfig{}, true) + + http.HandleFunc("/run_live", func(w http.ResponseWriter, req *http.Request) { + err := controller.RunLiveHandler(w, req) + if err != nil { + log.Printf("RunLiveHandler failed: %v", err) + } + }) + + fmt.Println("Serving UI on http://localhost:8081") + log.Fatal(http.ListenAndServe(":8081", nil)) +} diff --git a/examples/live/main.go b/examples/live/main.go index 4ec21cb41..025c255a1 100644 --- a/examples/live/main.go +++ b/examples/live/main.go @@ -71,7 +71,7 @@ func main() { ResponseModalities: []genai.Modality{genai.ModalityAudio}, OutputAudioTranscription: true, } - for ev, err := range r.RunLive(ctx, "user1", "sess1", queue, cfg) { + for ev, err := range r.RunLiveQueue(ctx, "user1", "sess1", queue, cfg) { if err != nil { fmt.Printf("[error] %v\n", err) break diff --git a/examples/mcp/main.go b/examples/mcp/main.go index 61de3e7cc..49c14362d 100644 --- a/examples/mcp/main.go +++ b/examples/mcp/main.go @@ -90,7 +90,7 @@ func main() { ctx, stop := signal.NotifyContext(context.Background(), os.Interrupt) defer stop() - model, err := gemini.NewModel(ctx, "gemini-2.5-flash", &genai.ClientConfig{ + model, err := gemini.NewModel(ctx, "gemini-3.1-flash-lite", &genai.ClientConfig{ APIKey: os.Getenv("GOOGLE_API_KEY"), }) if err != nil { diff --git a/examples/quickstart/main.go b/examples/quickstart/main.go index 8c1be10ad..85fd3c187 100644 --- a/examples/quickstart/main.go +++ b/examples/quickstart/main.go @@ -34,7 +34,7 @@ import ( func main() { ctx := context.Background() - model, err := gemini.NewModel(ctx, "gemini-2.5-flash", &genai.ClientConfig{ + model, err := gemini.NewModel(ctx, "gemini-3.1-flash-lite", &genai.ClientConfig{ APIKey: os.Getenv("GOOGLE_API_KEY"), }) if err != nil { diff --git a/examples/rest/main.go b/examples/rest/main.go index 84b57d2a6..e7fa1f859 100644 --- a/examples/rest/main.go +++ b/examples/rest/main.go @@ -37,7 +37,7 @@ func main() { ctx := context.Background() // Create a Gemini model - model, err := gemini.NewModel(ctx, "gemini-2.5-flash", &genai.ClientConfig{ + model, err := gemini.NewModel(ctx, "gemini-3.1-flash-lite", &genai.ClientConfig{ APIKey: os.Getenv("GOOGLE_API_KEY"), }) if err != nil { diff --git a/examples/skills/main.go b/examples/skills/main.go index f5980046c..9841ff33d 100644 --- a/examples/skills/main.go +++ b/examples/skills/main.go @@ -47,7 +47,7 @@ import ( func main() { ctx := context.Background() - model, err := gemini.NewModel(ctx, "gemini-2.5-flash", &genai.ClientConfig{ + model, err := gemini.NewModel(ctx, "gemini-3.1-flash-lite", &genai.ClientConfig{ APIKey: os.Getenv("GOOGLE_API_KEY"), }) if err != nil { diff --git a/examples/telemetry/main.go b/examples/telemetry/main.go index 88bdb4880..698a9b1d2 100644 --- a/examples/telemetry/main.go +++ b/examples/telemetry/main.go @@ -44,7 +44,7 @@ func main() { func run() error { ctx := context.Background() - model, err := gemini.NewModel(ctx, "gemini-2.5-flash", &genai.ClientConfig{ + model, err := gemini.NewModel(ctx, "gemini-3.1-flash-lite", &genai.ClientConfig{ APIKey: os.Getenv("GOOGLE_API_KEY"), }) if err != nil { diff --git a/examples/toolconfirmation/main.go b/examples/toolconfirmation/main.go index 16f49c8f7..c1c498210 100644 --- a/examples/toolconfirmation/main.go +++ b/examples/toolconfirmation/main.go @@ -79,7 +79,7 @@ var ( func main() { ctx := context.Background() - model, err := gemini.NewModel(ctx, "gemini-2.5-flash", &genai.ClientConfig{}) + model, err := gemini.NewModel(ctx, "gemini-3.1-flash-lite", &genai.ClientConfig{}) if err != nil { log.Fatalf("Failed to create model: %v", err) } @@ -126,7 +126,7 @@ func main() { } // requestVacationDays simulates the *initiation* of a long-running ticket creation task. -func requestVacationDays(ctx tool.Context, args RequestVacationArgs) (*RequestVacationResults, error) { +func requestVacationDays(ctx agent.ToolContext, args RequestVacationArgs) (*RequestVacationResults, error) { log.Printf("TOOL_EXEC: 'requestVacationDays' called with days: %d for user %s (Call ID: %s)\n", args.Days, args.UserID, ctx.FunctionCallID()) if args.Days <= 0 { diff --git a/examples/tools/loadartifacts/main.go b/examples/tools/loadartifacts/main.go index 4f55fb33b..debb23e1e 100644 --- a/examples/tools/loadartifacts/main.go +++ b/examples/tools/loadartifacts/main.go @@ -39,7 +39,7 @@ import ( func main() { ctx := context.Background() - model, err := gemini.NewModel(ctx, "gemini-2.5-flash", &genai.ClientConfig{ + model, err := gemini.NewModel(ctx, "gemini-3.1-flash-lite", &genai.ClientConfig{ APIKey: os.Getenv("GOOGLE_API_KEY"), }) if err != nil { diff --git a/examples/tools/loadmemory/main.go b/examples/tools/loadmemory/main.go index c7e42b1d3..41d192c1d 100644 --- a/examples/tools/loadmemory/main.go +++ b/examples/tools/loadmemory/main.go @@ -40,7 +40,7 @@ import ( func main() { ctx := context.Background() - model, err := gemini.NewModel(ctx, "gemini-2.5-flash", &genai.ClientConfig{ + model, err := gemini.NewModel(ctx, "gemini-3.1-flash-lite", &genai.ClientConfig{ APIKey: os.Getenv("GOOGLE_API_KEY"), }) if err != nil { @@ -170,7 +170,7 @@ func createPreviousSessionWithHistory( } for _, e := range events { - event := session.NewEvent("previous-session") + event := session.NewEventWithContext(ctx, "previous-session") event.Author = e.author event.LLMResponse = model.LLMResponse{ Content: genai.NewContentFromText(e.content, genai.Role(e.author)), diff --git a/examples/tools/multipletools/main.go b/examples/tools/multipletools/main.go index 993647b39..406e4b578 100644 --- a/examples/tools/multipletools/main.go +++ b/examples/tools/multipletools/main.go @@ -42,7 +42,7 @@ import ( func main() { ctx := context.Background() - model, err := gemini.NewModel(ctx, "gemini-2.5-flash", &genai.ClientConfig{ + model, err := gemini.NewModel(ctx, "gemini-3.1-flash-lite", &genai.ClientConfig{ APIKey: os.Getenv("GOOGLE_API_KEY"), }) if err != nil { @@ -68,7 +68,7 @@ func main() { type Output struct { Poem string `json:"poem"` } - handler := func(ctx tool.Context, input Input) (Output, error) { + handler := func(ctx agent.ToolContext, input Input) (Output, error) { return Output{ Poem: strings.Repeat("A line of a poem,", input.LineCount) + "\n", }, nil diff --git a/examples/vertexai/imagegenerator/main.go b/examples/vertexai/imagegenerator/main.go index 15efc7eb7..ec8e4bfb7 100644 --- a/examples/vertexai/imagegenerator/main.go +++ b/examples/vertexai/imagegenerator/main.go @@ -39,7 +39,7 @@ import ( func main() { ctx := context.Background() - model, err := gemini.NewModel(ctx, "gemini-2.0-flash-001", nil) + model, err := gemini.NewModel(ctx, "gemini-3.1-flash-image-preview", nil) if err != nil { log.Fatalf("Failed to create model: %v", err) } @@ -91,7 +91,7 @@ func main() { } // This is a function tool to generate images using Vertex AI's Imagen model. -func generateImage(ctx tool.Context, input generateImageInput) (generateImageResult, error) { +func generateImage(ctx agent.ToolContext, input generateImageInput) (generateImageResult, error) { client, err := genai.NewClient(ctx, &genai.ClientConfig{ Project: os.Getenv("GOOGLE_CLOUD_PROJECT"), Location: os.Getenv("GOOGLE_CLOUD_LOCATION"), @@ -129,7 +129,7 @@ type generateImageResult struct { // This is function tool that loads image from the artifacts service and // saves is to the local filesystem. -func saveImage(ctx tool.Context, input saveImageInput) (saveImageResult, error) { +func saveImage(ctx agent.ToolContext, input saveImageInput) (saveImageResult, error) { filename := input.Filename resp, err := ctx.Artifacts().Load(ctx, filename) if err != nil { diff --git a/examples/web/agents/image_generator.go b/examples/web/agents/image_generator.go index a325cdd6e..ffb1d76ac 100644 --- a/examples/web/agents/image_generator.go +++ b/examples/web/agents/image_generator.go @@ -30,7 +30,7 @@ import ( "google.golang.org/adk/tool/loadartifactstool" ) -func generateImage(ctx tool.Context, input generateImageInput) (generateImageResult, error) { +func generateImage(ctx agent.ToolContext, input generateImageInput) (generateImageResult, error) { client, err := genai.NewClient(ctx, &genai.ClientConfig{ Project: os.Getenv("GOOGLE_CLOUD_PROJECT"), Location: os.Getenv("GOOGLE_CLOUD_LOCATION"), diff --git a/examples/web/main.go b/examples/web/main.go index 6af33f4d4..de6bd3960 100644 --- a/examples/web/main.go +++ b/examples/web/main.go @@ -19,7 +19,7 @@ import ( "log" "os" - "github.com/a2aproject/a2a-go/a2asrv" + "github.com/a2aproject/a2a-go/v2/a2asrv" "github.com/google/uuid" "google.golang.org/genai" @@ -55,18 +55,16 @@ type AuthInterceptor struct { } // Before implements a before request callback. -func (a *AuthInterceptor) Before(ctx context.Context, callCtx *a2asrv.CallContext, req *a2asrv.Request) (context.Context, error) { - callCtx.User = &a2asrv.AuthenticatedUser{ - UserName: "user", - } - return ctx, nil +func (a *AuthInterceptor) Before(ctx context.Context, callCtx *a2asrv.CallContext, req *a2asrv.Request) (context.Context, any, error) { + callCtx.User = a2asrv.NewAuthenticatedUser("user", nil) + return ctx, nil, nil } func main() { ctx := context.Background() apiKey := os.Getenv("GOOGLE_API_KEY") - model, err := gemini.NewModel(ctx, "gemini-2.5-flash", &genai.ClientConfig{ + model, err := gemini.NewModel(ctx, "gemini-3.1-flash-lite", &genai.ClientConfig{ APIKey: apiKey, }) if err != nil { @@ -105,7 +103,7 @@ func main() { SessionService: sessionService, AgentLoader: agentLoader, A2AOptions: []a2asrv.RequestHandlerOption{ - a2asrv.WithCallInterceptor(&AuthInterceptor{}), + a2asrv.WithCallInterceptors(&AuthInterceptor{}), }, } diff --git a/examples/workflowagents/sequentialCode/main.go b/examples/workflowagents/sequentialCode/main.go index 70f3bf9ca..7ecc98639 100644 --- a/examples/workflowagents/sequentialCode/main.go +++ b/examples/workflowagents/sequentialCode/main.go @@ -33,7 +33,7 @@ import ( func main() { ctx := context.Background() - model, err := gemini.NewModel(ctx, "gemini-2.5-flash", &genai.ClientConfig{}) + model, err := gemini.NewModel(ctx, "gemini-3.1-flash-lite", &genai.ClientConfig{}) if err != nil { log.Fatalf("failed to create model: %s", err) } diff --git a/go.mod b/go.mod index 68e5e1b26..2f606175a 100644 --- a/go.mod +++ b/go.mod @@ -6,7 +6,7 @@ require ( cloud.google.com/go v0.123.0 cloud.google.com/go/aiplatform v1.121.0 cloud.google.com/go/storage v1.56.1 - github.com/a2aproject/a2a-go v0.3.13 + github.com/a2aproject/a2a-go v0.3.15 github.com/awalterschulze/gographviz v2.0.3+incompatible github.com/glebarez/sqlite v1.8.0 github.com/google/go-cmp v0.7.0 @@ -16,19 +16,19 @@ require ( github.com/gorilla/mux v1.8.1 github.com/mitchellh/mapstructure v1.5.0 github.com/modelcontextprotocol/go-sdk v1.4.1 - github.com/spf13/cobra v1.8.1 - go.opentelemetry.io/contrib/detectors/gcp v1.40.0 - go.opentelemetry.io/otel v1.40.0 - go.opentelemetry.io/otel/exporters/otlp/otlplog/otlploghttp v0.16.0 - go.opentelemetry.io/otel/exporters/otlp/otlptrace/otlptracehttp v1.39.0 - go.opentelemetry.io/otel/log v0.16.0 - go.opentelemetry.io/otel/sdk v1.40.0 - go.opentelemetry.io/otel/trace v1.40.0 + github.com/spf13/cobra v1.10.2 + go.opentelemetry.io/contrib/detectors/gcp v1.42.0 + go.opentelemetry.io/otel v1.43.0 + go.opentelemetry.io/otel/exporters/otlp/otlplog/otlploghttp v0.19.0 + go.opentelemetry.io/otel/exporters/otlp/otlptrace/otlptracehttp v1.43.0 + go.opentelemetry.io/otel/log v0.19.0 + go.opentelemetry.io/otel/sdk v1.43.0 + go.opentelemetry.io/otel/trace v1.43.0 golang.org/x/oauth2 v0.36.0 golang.org/x/sync v0.20.0 - google.golang.org/api v0.272.0 - google.golang.org/genai v1.40.0 - google.golang.org/grpc v1.79.3 + google.golang.org/api v0.279.0 + google.golang.org/genai v1.57.0 + google.golang.org/grpc v1.81.0 google.golang.org/protobuf v1.36.11 gopkg.in/yaml.v3 v3.0.1 gorm.io/gorm v1.31.0 @@ -38,9 +38,14 @@ require ( require github.com/hashicorp/golang-lru/v2 v2.0.7 +require ( + github.com/a2aproject/a2a-go/v2 v2.3.1 + golang.org/x/mod v0.35.0 // indirect +) + require ( cel.dev/expr v0.25.1 // indirect - cloud.google.com/go/auth v0.18.2 // indirect + cloud.google.com/go/auth v0.20.0 // indirect cloud.google.com/go/auth/oauth2adapt v0.2.8 // indirect cloud.google.com/go/compute/metadata v0.9.0 // indirect cloud.google.com/go/iam v1.5.3 // indirect @@ -51,20 +56,20 @@ require ( github.com/GoogleCloudPlatform/opentelemetry-operations-go/internal/resourcemapping v0.53.0 // indirect github.com/cenkalti/backoff/v5 v5.0.3 // indirect github.com/cespare/xxhash/v2 v2.3.0 // indirect - github.com/cncf/xds/go v0.0.0-20251210132809-ee656c7534f5 // indirect + github.com/cncf/xds/go v0.0.0-20260202195803-dba9d589def2 // indirect github.com/dustin/go-humanize v1.0.1 // indirect - github.com/envoyproxy/go-control-plane/envoy v1.36.0 // indirect - github.com/envoyproxy/protoc-gen-validate v1.3.0 // indirect + github.com/envoyproxy/go-control-plane/envoy v1.37.0 // indirect + github.com/envoyproxy/protoc-gen-validate v1.3.3 // indirect github.com/felixge/httpsnoop v1.0.4 // indirect github.com/glebarez/go-sqlite v1.21.1 // indirect - github.com/go-jose/go-jose/v4 v4.1.3 // indirect + github.com/go-jose/go-jose/v4 v4.1.4 // indirect github.com/go-logr/logr v1.4.3 // indirect github.com/go-logr/stdr v1.2.2 // indirect github.com/google/s2a-go v0.1.9 // indirect - github.com/googleapis/enterprise-certificate-proxy v0.3.14 // indirect - github.com/googleapis/gax-go/v2 v2.18.0 // indirect - github.com/gorilla/websocket v1.5.3 // indirect - github.com/grpc-ecosystem/grpc-gateway/v2 v2.27.7 // indirect + github.com/googleapis/enterprise-certificate-proxy v0.3.15 // indirect + github.com/googleapis/gax-go/v2 v2.22.0 // indirect + github.com/gorilla/websocket v1.5.3 + github.com/grpc-ecosystem/grpc-gateway/v2 v2.28.0 // indirect github.com/inconshreveable/mousetrap v1.1.0 // indirect github.com/jinzhu/inflection v1.0.0 // indirect github.com/jinzhu/now v1.1.5 // indirect @@ -77,21 +82,21 @@ require ( github.com/spiffe/go-spiffe/v2 v2.6.0 // indirect github.com/yosida95/uritemplate/v3 v3.0.2 // indirect go.opentelemetry.io/auto/sdk v1.2.1 // indirect - go.opentelemetry.io/contrib/instrumentation/google.golang.org/grpc/otelgrpc v0.63.0 // indirect - go.opentelemetry.io/contrib/instrumentation/net/http/otelhttp v0.63.0 // indirect - go.opentelemetry.io/otel/exporters/otlp/otlptrace v1.39.0 // indirect - go.opentelemetry.io/otel/metric v1.40.0 // indirect - go.opentelemetry.io/otel/sdk/log v0.16.0 - go.opentelemetry.io/otel/sdk/metric v1.40.0 // indirect - go.opentelemetry.io/proto/otlp v1.9.0 // indirect - golang.org/x/crypto v0.49.0 // indirect - golang.org/x/net v0.52.0 // indirect - golang.org/x/sys v0.42.0 // indirect - golang.org/x/text v0.35.0 // indirect + go.opentelemetry.io/contrib/instrumentation/google.golang.org/grpc/otelgrpc v0.67.0 // indirect + go.opentelemetry.io/contrib/instrumentation/net/http/otelhttp v0.68.0 // indirect + go.opentelemetry.io/otel/exporters/otlp/otlptrace v1.43.0 // indirect + go.opentelemetry.io/otel/metric v1.43.0 // indirect + go.opentelemetry.io/otel/sdk/log v0.19.0 + go.opentelemetry.io/otel/sdk/metric v1.43.0 // indirect + go.opentelemetry.io/proto/otlp v1.10.0 // indirect + golang.org/x/crypto v0.51.0 // indirect + golang.org/x/net v0.55.0 // indirect + golang.org/x/sys v0.45.0 // indirect + golang.org/x/text v0.37.0 // indirect golang.org/x/time v0.15.0 // indirect - google.golang.org/genproto v0.0.0-20260217215200-42d3e9bedb6d // indirect - google.golang.org/genproto/googleapis/api v0.0.0-20260316180232-0b37fe3546d5 // indirect - google.golang.org/genproto/googleapis/rpc v0.0.0-20260316180232-0b37fe3546d5 // indirect + google.golang.org/genproto v0.0.0-20260319201613-d00831a3d3e7 // indirect + google.golang.org/genproto/googleapis/api v0.0.0-20260427160629-7cedc36a6bc4 // indirect + google.golang.org/genproto/googleapis/rpc v0.0.0-20260511170946-3700d4141b60 // indirect modernc.org/libc v1.22.3 // indirect modernc.org/mathutil v1.5.0 // indirect modernc.org/memory v1.5.0 // indirect diff --git a/go.sum b/go.sum index 2ac6ef276..46698917e 100644 --- a/go.sum +++ b/go.sum @@ -4,8 +4,8 @@ cloud.google.com/go v0.123.0 h1:2NAUJwPR47q+E35uaJeYoNhuNEM9kM8SjgRgdeOJUSE= cloud.google.com/go v0.123.0/go.mod h1:xBoMV08QcqUGuPW65Qfm1o9Y4zKZBpGS+7bImXLTAZU= cloud.google.com/go/aiplatform v1.121.0 h1:8y8sNfVAW1DVhFbSbI7d8rrqBGGJFk6EoV6atidlyQc= cloud.google.com/go/aiplatform v1.121.0/go.mod h1:juMdDWeNphHV40KhWdN+563zNCOKNmLJjk5D2TA43ls= -cloud.google.com/go/auth v0.18.2 h1:+Nbt5Ev0xEqxlNjd6c+yYUeosQ5TtEUaNcN/3FozlaM= -cloud.google.com/go/auth v0.18.2/go.mod h1:xD+oY7gcahcu7G2SG2DsBerfFxgPAJz17zz2joOFF3M= +cloud.google.com/go/auth v0.20.0 h1:kXTssoVb4azsVDoUiF8KvxAqrsQcQtB53DcSgta74CA= +cloud.google.com/go/auth v0.20.0/go.mod h1:942/yi/itH1SsmpyrbnTMDgGfdy2BUqIKyd0cyYLc5Q= cloud.google.com/go/auth/oauth2adapt v0.2.8 h1:keo8NaayQZ6wimpNSmW5OPc283g65QNIiLpZnkHRbnc= cloud.google.com/go/auth/oauth2adapt v0.2.8/go.mod h1:XQ9y31RkqZCcwJWNSx2Xvric3RrU88hAYYbjDWYDL+c= cloud.google.com/go/compute/metadata v0.9.0 h1:pDUj4QMoPejqq20dK0Pg2N4yG9zIkYGdBtwLoEkH9Zs= @@ -30,37 +30,39 @@ github.com/GoogleCloudPlatform/opentelemetry-operations-go/internal/cloudmock v0 github.com/GoogleCloudPlatform/opentelemetry-operations-go/internal/cloudmock v0.53.0/go.mod h1:jUZ5LYlw40WMd07qxcQJD5M40aUxrfwqQX1g7zxYnrQ= github.com/GoogleCloudPlatform/opentelemetry-operations-go/internal/resourcemapping v0.53.0 h1:Ron4zCA/yk6U7WOBXhTJcDpsUBG9npumK6xw2auFltQ= github.com/GoogleCloudPlatform/opentelemetry-operations-go/internal/resourcemapping v0.53.0/go.mod h1:cSgYe11MCNYunTnRXrKiR/tHc0eoKjICUuWpNZoVCOo= -github.com/a2aproject/a2a-go v0.3.13 h1:WpIcSHgCySIxD7OQEdV7U7WJc/HL/G2QQj0RJ0YhPi0= -github.com/a2aproject/a2a-go v0.3.13/go.mod h1:I7Cm+a1oL+UT6zMoP+roaRE5vdfUa1iQGVN8aSOuZ0I= +github.com/a2aproject/a2a-go v0.3.15 h1:h5YpCiPq3jxQ5rIns7oDjPag3ivP8u817AzdA4F+NiI= +github.com/a2aproject/a2a-go v0.3.15/go.mod h1:I7Cm+a1oL+UT6zMoP+roaRE5vdfUa1iQGVN8aSOuZ0I= +github.com/a2aproject/a2a-go/v2 v2.3.1 h1:QWMdOX2UsJ8BJmjs952eo1FRyGsOVl0gFCKeM76AgGE= +github.com/a2aproject/a2a-go/v2 v2.3.1/go.mod h1:mkZr8y2bUgAVQsjs/5fHK7xrRlAHDybMEyxWh2tKRC8= github.com/awalterschulze/gographviz v2.0.3+incompatible h1:9sVEXJBJLwGX7EQVhLm2elIKCm7P2YHFC8v6096G09E= github.com/awalterschulze/gographviz v2.0.3+incompatible/go.mod h1:GEV5wmg4YquNw7v1kkyoX9etIk8yVmXj+AkDHuuETHs= github.com/cenkalti/backoff/v5 v5.0.3 h1:ZN+IMa753KfX5hd8vVaMixjnqRZ3y8CuJKRKj1xcsSM= github.com/cenkalti/backoff/v5 v5.0.3/go.mod h1:rkhZdG3JZukswDf7f0cwqPNk4K0sa+F97BxZthm/crw= github.com/cespare/xxhash/v2 v2.3.0 h1:UL815xU9SqsFlibzuggzjXhog7bL6oX9BbNZnL2UFvs= github.com/cespare/xxhash/v2 v2.3.0/go.mod h1:VGX0DQ3Q6kWi7AoAeZDth3/j3BFtOZR5XLFGgcrjCOs= -github.com/cncf/xds/go v0.0.0-20251210132809-ee656c7534f5 h1:6xNmx7iTtyBRev0+D/Tv1FZd4SCg8axKApyNyRsAt/w= -github.com/cncf/xds/go v0.0.0-20251210132809-ee656c7534f5/go.mod h1:KdCmV+x/BuvyMxRnYBlmVaq4OLiKW6iRQfvC62cvdkI= -github.com/cpuguy83/go-md2man/v2 v2.0.4/go.mod h1:tgQtvFlXSQOSOSIRvRPT7W67SCa46tRHOmNcaadrF8o= +github.com/cncf/xds/go v0.0.0-20260202195803-dba9d589def2 h1:aBangftG7EVZoUb69Os8IaYg++6uMOdKK83QtkkvJik= +github.com/cncf/xds/go v0.0.0-20260202195803-dba9d589def2/go.mod h1:qwXFYgsP6T7XnJtbKlf1HP8AjxZZyzxMmc+Lq5GjlU4= +github.com/cpuguy83/go-md2man/v2 v2.0.6/go.mod h1:oOW0eioCTA6cOiMLiUPZOpcVxMig6NIQQ7OS05n1F4g= github.com/davecgh/go-spew v1.1.2-0.20180830191138-d8f796af33cc h1:U9qPSI2PIWSS1VwoXQT9A3Wy9MM3WgvqSxFWenqJduM= github.com/davecgh/go-spew v1.1.2-0.20180830191138-d8f796af33cc/go.mod h1:J7Y8YcW2NihsgmVo/mv3lAwl/skON4iLHjSsI+c5H38= github.com/dustin/go-humanize v1.0.1 h1:GzkhY7T5VNhEkwH0PVJgjz+fX1rhBrR7pRT3mDkpeCY= github.com/dustin/go-humanize v1.0.1/go.mod h1:Mu1zIs6XwVuF/gI1OepvI0qD18qycQx+mFykh5fBlto= github.com/envoyproxy/go-control-plane v0.14.0 h1:hbG2kr4RuFj222B6+7T83thSPqLjwBIfQawTkC++2HA= github.com/envoyproxy/go-control-plane v0.14.0/go.mod h1:NcS5X47pLl/hfqxU70yPwL9ZMkUlwlKxtAohpi2wBEU= -github.com/envoyproxy/go-control-plane/envoy v1.36.0 h1:yg/JjO5E7ubRyKX3m07GF3reDNEnfOboJ0QySbH736g= -github.com/envoyproxy/go-control-plane/envoy v1.36.0/go.mod h1:ty89S1YCCVruQAm9OtKeEkQLTb+Lkz0k8v9W0Oxsv98= +github.com/envoyproxy/go-control-plane/envoy v1.37.0 h1:u3riX6BoYRfF4Dr7dwSOroNfdSbEPe9Yyl09/B6wBrQ= +github.com/envoyproxy/go-control-plane/envoy v1.37.0/go.mod h1:DReE9MMrmecPy+YvQOAOHNYMALuowAnbjjEMkkWOi6A= github.com/envoyproxy/go-control-plane/ratelimit v0.1.0 h1:/G9QYbddjL25KvtKTv3an9lx6VBE2cnb8wp1vEGNYGI= github.com/envoyproxy/go-control-plane/ratelimit v0.1.0/go.mod h1:Wk+tMFAFbCXaJPzVVHnPgRKdUdwW/KdbRt94AzgRee4= -github.com/envoyproxy/protoc-gen-validate v1.3.0 h1:TvGH1wof4H33rezVKWSpqKz5NXWg5VPuZ0uONDT6eb4= -github.com/envoyproxy/protoc-gen-validate v1.3.0/go.mod h1:HvYl7zwPa5mffgyeTUHA9zHIH36nmrm7oCbo4YKoSWA= +github.com/envoyproxy/protoc-gen-validate v1.3.3 h1:MVQghNeW+LZcmXe7SY1V36Z+WFMDjpqGAGacLe2T0ds= +github.com/envoyproxy/protoc-gen-validate v1.3.3/go.mod h1:TsndJ/ngyIdQRhMcVVGDDHINPLWB7C82oDArY51KfB0= github.com/felixge/httpsnoop v1.0.4 h1:NFTV2Zj1bL4mc9sqWACXbQFVBBg2W3GPvqp8/ESS2Wg= github.com/felixge/httpsnoop v1.0.4/go.mod h1:m8KPJKqk1gH5J9DgRY2ASl2lWCfGKXixSwevea8zH2U= github.com/glebarez/go-sqlite v1.21.1 h1:7MZyUPh2XTrHS7xNEHQbrhfMZuPSzhkm2A1qgg0y5NY= github.com/glebarez/go-sqlite v1.21.1/go.mod h1:ISs8MF6yk5cL4n/43rSOmVMGJJjHYr7L2MbZZ5Q4E2E= github.com/glebarez/sqlite v1.8.0 h1:02X12E2I/4C1n+v90yTqrjRa8yuo7c3KeHI3FRznCvc= github.com/glebarez/sqlite v1.8.0/go.mod h1:bpET16h1za2KOOMb8+jCp6UBP/iahDpfPQqSaYLTLx8= -github.com/go-jose/go-jose/v4 v4.1.3 h1:CVLmWDhDVRa6Mi/IgCgaopNosCaHz7zrMeF9MlZRkrs= -github.com/go-jose/go-jose/v4 v4.1.3/go.mod h1:x4oUasVrzR7071A4TnHLGSPpNOm2a21K9Kf04k1rs08= +github.com/go-jose/go-jose/v4 v4.1.4 h1:moDMcTHmvE6Groj34emNPLs/qtYXRVcd6S7NHbHz3kA= +github.com/go-jose/go-jose/v4 v4.1.4/go.mod h1:x4oUasVrzR7071A4TnHLGSPpNOm2a21K9Kf04k1rs08= github.com/go-logr/logr v1.2.2/go.mod h1:jdQByPbusPIv2/zmleS9BjJVeZ6kBagPoEUsqbVz/1A= github.com/go-logr/logr v1.4.3 h1:CjnDlHq8ikf6E492q6eKboGOC0T8CDaOvkHCIg8idEI= github.com/go-logr/logr v1.4.3/go.mod h1:9T104GzyrTigFIr8wt5mBrctHMim0Nb2HLGrmQ40KvY= @@ -84,16 +86,16 @@ github.com/google/safehtml v0.1.0 h1:EwLKo8qawTKfsi0orxcQAZzu07cICaBeFMegAU9eaT8 github.com/google/safehtml v0.1.0/go.mod h1:L4KWwDsUJdECRAEpZoBn3O64bQaywRscowZjJAzjHnU= github.com/google/uuid v1.6.0 h1:NIvaJDMOsjHA8n1jAhLSgzrAzy1Hgr+hNrb57e+94F0= github.com/google/uuid v1.6.0/go.mod h1:TIyPZe4MgqvfeYDBFedMoGGpEw/LqOeaOT+nhxU+yHo= -github.com/googleapis/enterprise-certificate-proxy v0.3.14 h1:yh8ncqsbUY4shRD5dA6RlzjJaT4hi3kII+zYw8wmLb8= -github.com/googleapis/enterprise-certificate-proxy v0.3.14/go.mod h1:vqVt9yG9480NtzREnTlmGSBmFrA+bzb0yl0TxoBQXOg= -github.com/googleapis/gax-go/v2 v2.18.0 h1:jxP5Uuo3bxm3M6gGtV94P4lliVetoCB4Wk2x8QA86LI= -github.com/googleapis/gax-go/v2 v2.18.0/go.mod h1:uSzZN4a356eRG985CzJ3WfbFSpqkLTjsnhWGJR6EwrE= +github.com/googleapis/enterprise-certificate-proxy v0.3.15 h1:xolVQTEXusUcAA5UgtyRLjelpFFHWlPQ4XfWGc7MBas= +github.com/googleapis/enterprise-certificate-proxy v0.3.15/go.mod h1:vqVt9yG9480NtzREnTlmGSBmFrA+bzb0yl0TxoBQXOg= +github.com/googleapis/gax-go/v2 v2.22.0 h1:PjIWBpgGIVKGoCXuiCoP64altEJCj3/Ei+kSU5vlZD4= +github.com/googleapis/gax-go/v2 v2.22.0/go.mod h1:irWBbALSr0Sk3qlqb9SyJ1h68WjgeFuiOzI4Rqw5+aY= github.com/gorilla/mux v1.8.1 h1:TuBL49tXwgrFYWhqrNgrUNEY92u81SPhu7sTdzQEiWY= github.com/gorilla/mux v1.8.1/go.mod h1:AKf9I4AEqPTmMytcMc0KkNouC66V3BtZ4qD5fmWSiMQ= github.com/gorilla/websocket v1.5.3 h1:saDtZ6Pbx/0u+bgYQ3q96pZgCzfhKXGPqt7kZ72aNNg= github.com/gorilla/websocket v1.5.3/go.mod h1:YR8l580nyteQvAITg2hZ9XVh4b55+EU/adAjf1fMHhE= -github.com/grpc-ecosystem/grpc-gateway/v2 v2.27.7 h1:X+2YciYSxvMQK0UZ7sg45ZVabVZBeBuvMkmuI2V3Fak= -github.com/grpc-ecosystem/grpc-gateway/v2 v2.27.7/go.mod h1:lW34nIZuQ8UDPdkon5fmfp2l3+ZkQ2me/+oecHYLOII= +github.com/grpc-ecosystem/grpc-gateway/v2 v2.28.0 h1:HWRh5R2+9EifMyIHV7ZV+MIZqgz+PMpZ14Jynv3O2Zs= +github.com/grpc-ecosystem/grpc-gateway/v2 v2.28.0/go.mod h1:JfhWUomR1baixubs02l85lZYYOm7LV6om4ceouMv45c= github.com/hashicorp/golang-lru/v2 v2.0.7 h1:a+bsQ5rvGLjzHuww6tVxozPZFVghXaHOwFs4luLUK2k= github.com/hashicorp/golang-lru/v2 v2.0.7/go.mod h1:QeFd9opnmA6QUJc5vARoKUSoFhyfM2/ZepoAG6RGpeM= github.com/inconshreveable/mousetrap v1.1.0 h1:wN+x4NVGpMsO7ErUn/mUI3vEoE6Jt13X2s0bqwp9tc8= @@ -126,9 +128,9 @@ github.com/segmentio/asm v1.1.3 h1:WM03sfUOENvvKexOLp+pCqgb/WDjsi7EK8gIsICtzhc= github.com/segmentio/asm v1.1.3/go.mod h1:Ld3L4ZXGNcSLRg4JBsZ3//1+f/TjYl0Mzen/DQy1EJg= github.com/segmentio/encoding v0.5.4 h1:OW1VRern8Nw6ITAtwSZ7Idrl3MXCFwXHPgqESYfvNt0= github.com/segmentio/encoding v0.5.4/go.mod h1:HS1ZKa3kSN32ZHVZ7ZLPLXWvOVIiZtyJnO1gPH1sKt0= -github.com/spf13/cobra v1.8.1 h1:e5/vxKd/rZsfSJMUX1agtjeTDf+qv1/JdBF8gg5k9ZM= -github.com/spf13/cobra v1.8.1/go.mod h1:wHxEcudfqmLYa8iTfL+OuZPbBZkmvliBWKIezN3kD9Y= -github.com/spf13/pflag v1.0.5/go.mod h1:McXfInJRrz4CZXVZOBLb0bTZqETkiAhM9Iw0y3An2Bg= +github.com/spf13/cobra v1.10.2 h1:DMTTonx5m65Ic0GOoRY2c16WCbHxOOw6xxezuLaBpcU= +github.com/spf13/cobra v1.10.2/go.mod h1:7C1pvHqHw5A4vrJfjNwvOdzYu0Gml16OCs2GRiTUUS4= +github.com/spf13/pflag v1.0.9/go.mod h1:McXfInJRrz4CZXVZOBLb0bTZqETkiAhM9Iw0y3An2Bg= github.com/spf13/pflag v1.0.10 h1:4EBh2KAYBwaONj6b2Ye1GiHfwjqyROoF4RwYO+vPwFk= github.com/spf13/pflag v1.0.10/go.mod h1:McXfInJRrz4CZXVZOBLb0bTZqETkiAhM9Iw0y3An2Bg= github.com/spiffe/go-spiffe/v2 v2.6.0 h1:l+DolpxNWYgruGQVV0xsfeya3CsC7m8iBzDnMpsbLuo= @@ -139,73 +141,76 @@ github.com/yosida95/uritemplate/v3 v3.0.2 h1:Ed3Oyj9yrmi9087+NczuL5BwkIc4wvTb5zI github.com/yosida95/uritemplate/v3 v3.0.2/go.mod h1:ILOh0sOhIJR3+L/8afwt/kE++YT040gmv5BQTMR2HP4= go.opentelemetry.io/auto/sdk v1.2.1 h1:jXsnJ4Lmnqd11kwkBV2LgLoFMZKizbCi5fNZ/ipaZ64= go.opentelemetry.io/auto/sdk v1.2.1/go.mod h1:KRTj+aOaElaLi+wW1kO/DZRXwkF4C5xPbEe3ZiIhN7Y= -go.opentelemetry.io/contrib/detectors/gcp v1.40.0 h1:Awaf8gmW99tZTOWqkLCOl6aw1/rxAWVlHsHIZ3fT2sA= -go.opentelemetry.io/contrib/detectors/gcp v1.40.0/go.mod h1:99OY9ZCqyLkzJLTh5XhECpLRSxcZl+ZDKBEO+jMBFR4= -go.opentelemetry.io/contrib/instrumentation/google.golang.org/grpc/otelgrpc v0.63.0 h1:YH4g8lQroajqUwWbq/tr2QX1JFmEXaDLgG+ew9bLMWo= -go.opentelemetry.io/contrib/instrumentation/google.golang.org/grpc/otelgrpc v0.63.0/go.mod h1:fvPi2qXDqFs8M4B4fmJhE92TyQs9Ydjlg3RvfUp+NbQ= -go.opentelemetry.io/contrib/instrumentation/net/http/otelhttp v0.63.0 h1:RbKq8BG0FI8OiXhBfcRtqqHcZcka+gU3cskNuf05R18= -go.opentelemetry.io/contrib/instrumentation/net/http/otelhttp v0.63.0/go.mod h1:h06DGIukJOevXaj/xrNjhi/2098RZzcLTbc0jDAUbsg= -go.opentelemetry.io/otel v1.40.0 h1:oA5YeOcpRTXq6NN7frwmwFR0Cn3RhTVZvXsP4duvCms= -go.opentelemetry.io/otel v1.40.0/go.mod h1:IMb+uXZUKkMXdPddhwAHm6UfOwJyh4ct1ybIlV14J0g= -go.opentelemetry.io/otel/exporters/otlp/otlplog/otlploghttp v0.16.0 h1:djrxvDxAe44mJUrKataUbOhCKhR3F8QCyWucO16hTQs= -go.opentelemetry.io/otel/exporters/otlp/otlplog/otlploghttp v0.16.0/go.mod h1:dt3nxpQEiSoKvfTVxp3TUg5fHPLhKtbcnN3Z1I1ePD0= -go.opentelemetry.io/otel/exporters/otlp/otlptrace v1.39.0 h1:f0cb2XPmrqn4XMy9PNliTgRKJgS5WcL/u0/WRYGz4t0= -go.opentelemetry.io/otel/exporters/otlp/otlptrace v1.39.0/go.mod h1:vnakAaFckOMiMtOIhFI2MNH4FYrZzXCYxmb1LlhoGz8= -go.opentelemetry.io/otel/exporters/otlp/otlptrace/otlptracehttp v1.39.0 h1:Ckwye2FpXkYgiHX7fyVrN1uA/UYd9ounqqTuSNAv0k4= -go.opentelemetry.io/otel/exporters/otlp/otlptrace/otlptracehttp v1.39.0/go.mod h1:teIFJh5pW2y+AN7riv6IBPX2DuesS3HgP39mwOspKwU= +go.opentelemetry.io/contrib/detectors/gcp v1.42.0 h1:kpt2PEJuOuqYkPcktfJqWWDjTEd/FNgrxcniL7kQrXQ= +go.opentelemetry.io/contrib/detectors/gcp v1.42.0/go.mod h1:W9zQ439utxymRrXsUOzZbFX4JhLxXU4+ZnCt8GG7yA8= +go.opentelemetry.io/contrib/instrumentation/google.golang.org/grpc/otelgrpc v0.67.0 h1:yI1/OhfEPy7J9eoa6Sj051C7n5dvpj0QX8g4sRchg04= +go.opentelemetry.io/contrib/instrumentation/google.golang.org/grpc/otelgrpc v0.67.0/go.mod h1:NoUCKYWK+3ecatC4HjkRktREheMeEtrXoQxrqYFeHSc= +go.opentelemetry.io/contrib/instrumentation/net/http/otelhttp v0.68.0 h1:CqXxU8VOmDefoh0+ztfGaymYbhdB/tT3zs79QaZTNGY= +go.opentelemetry.io/contrib/instrumentation/net/http/otelhttp v0.68.0/go.mod h1:BuhAPThV8PBHBvg8ZzZ/Ok3idOdhWIodywz2xEcRbJo= +go.opentelemetry.io/otel v1.43.0 h1:mYIM03dnh5zfN7HautFE4ieIig9amkNANT+xcVxAj9I= +go.opentelemetry.io/otel v1.43.0/go.mod h1:JuG+u74mvjvcm8vj8pI5XiHy1zDeoCS2LB1spIq7Ay0= +go.opentelemetry.io/otel/exporters/otlp/otlplog/otlploghttp v0.19.0 h1:HIBTQ3VO5aupLKjC90JgMqpezVXwFuq6Ryjn0/izoag= +go.opentelemetry.io/otel/exporters/otlp/otlplog/otlploghttp v0.19.0/go.mod h1:ji9vId85hMxqfvICA0Jt8JqEdrXaAkcpkI9HPXya0ro= +go.opentelemetry.io/otel/exporters/otlp/otlptrace v1.43.0 h1:88Y4s2C8oTui1LGM6bTWkw0ICGcOLCAI5l6zsD1j20k= +go.opentelemetry.io/otel/exporters/otlp/otlptrace v1.43.0/go.mod h1:Vl1/iaggsuRlrHf/hfPJPvVag77kKyvrLeD10kpMl+A= +go.opentelemetry.io/otel/exporters/otlp/otlptrace/otlptracehttp v1.43.0 h1:3iZJKlCZufyRzPzlQhUIWVmfltrXuGyfjREgGP3UUjc= +go.opentelemetry.io/otel/exporters/otlp/otlptrace/otlptracehttp v1.43.0/go.mod h1:/G+nUPfhq2e+qiXMGxMwumDrP5jtzU+mWN7/sjT2rak= go.opentelemetry.io/otel/exporters/stdout/stdoutmetric v1.36.0 h1:rixTyDGXFxRy1xzhKrotaHy3/KXdPhlWARrCgK+eqUY= go.opentelemetry.io/otel/exporters/stdout/stdoutmetric v1.36.0/go.mod h1:dowW6UsM9MKbJq5JTz2AMVp3/5iW5I/TStsk8S+CfHw= -go.opentelemetry.io/otel/log v0.16.0 h1:DeuBPqCi6pQwtCK0pO4fvMB5eBq6sNxEnuTs88pjsN4= -go.opentelemetry.io/otel/log v0.16.0/go.mod h1:rWsmqNVTLIA8UnwYVOItjyEZDbKIkMxdQunsIhpUMes= -go.opentelemetry.io/otel/metric v1.40.0 h1:rcZe317KPftE2rstWIBitCdVp89A2HqjkxR3c11+p9g= -go.opentelemetry.io/otel/metric v1.40.0/go.mod h1:ib/crwQH7N3r5kfiBZQbwrTge743UDc7DTFVZrrXnqc= -go.opentelemetry.io/otel/sdk v1.40.0 h1:KHW/jUzgo6wsPh9At46+h4upjtccTmuZCFAc9OJ71f8= -go.opentelemetry.io/otel/sdk v1.40.0/go.mod h1:Ph7EFdYvxq72Y8Li9q8KebuYUr2KoeyHx0DRMKrYBUE= -go.opentelemetry.io/otel/sdk/log v0.16.0 h1:e/b4bdlQwC5fnGtG3dlXUrNOnP7c8YLVSpSfEBIkTnI= -go.opentelemetry.io/otel/sdk/log v0.16.0/go.mod h1:JKfP3T6ycy7QEuv3Hj8oKDy7KItrEkus8XJE6EoSzw4= -go.opentelemetry.io/otel/sdk/log/logtest v0.16.0 h1:/XVkpZ41rVRTP4DfMgYv1nEtNmf65XPPyAdqV90TMy4= -go.opentelemetry.io/otel/sdk/log/logtest v0.16.0/go.mod h1:iOOPgQr5MY9oac/F5W86mXdeyWZGleIx3uXO98X2R6Y= -go.opentelemetry.io/otel/sdk/metric v1.40.0 h1:mtmdVqgQkeRxHgRv4qhyJduP3fYJRMX4AtAlbuWdCYw= -go.opentelemetry.io/otel/sdk/metric v1.40.0/go.mod h1:4Z2bGMf0KSK3uRjlczMOeMhKU2rhUqdWNoKcYrtcBPg= -go.opentelemetry.io/otel/trace v1.40.0 h1:WA4etStDttCSYuhwvEa8OP8I5EWu24lkOzp+ZYblVjw= -go.opentelemetry.io/otel/trace v1.40.0/go.mod h1:zeAhriXecNGP/s2SEG3+Y8X9ujcJOTqQ5RgdEJcawiA= -go.opentelemetry.io/proto/otlp v1.9.0 h1:l706jCMITVouPOqEnii2fIAuO3IVGBRPV5ICjceRb/A= -go.opentelemetry.io/proto/otlp v1.9.0/go.mod h1:xE+Cx5E/eEHw+ISFkwPLwCZefwVjY+pqKg1qcK03+/4= +go.opentelemetry.io/otel/log v0.19.0 h1:KUZs/GOsw79TBBMfDWsXS+KZ4g2Ckzksd1ymzsIEbo4= +go.opentelemetry.io/otel/log v0.19.0/go.mod h1:5DQYeGmxVIr4n0/BcJvF4upsraHjg6vudJJpnkL6Ipk= +go.opentelemetry.io/otel/metric v1.43.0 h1:d7638QeInOnuwOONPp4JAOGfbCEpYb+K6DVWvdxGzgM= +go.opentelemetry.io/otel/metric v1.43.0/go.mod h1:RDnPtIxvqlgO8GRW18W6Z/4P462ldprJtfxHxyKd2PY= +go.opentelemetry.io/otel/sdk v1.43.0 h1:pi5mE86i5rTeLXqoF/hhiBtUNcrAGHLKQdhg4h4V9Dg= +go.opentelemetry.io/otel/sdk v1.43.0/go.mod h1:P+IkVU3iWukmiit/Yf9AWvpyRDlUeBaRg6Y+C58QHzg= +go.opentelemetry.io/otel/sdk/log v0.19.0 h1:scYVLqT22D2gqXItnWiocLUKGH9yvkkeql5dBDiXyko= +go.opentelemetry.io/otel/sdk/log v0.19.0/go.mod h1:vFBowwXGLlW9AvpuF7bMgnNI95LiW10szrOdvzBHlAg= +go.opentelemetry.io/otel/sdk/log/logtest v0.19.0 h1:BEbF7ZBB6qQloV/Ub1+3NQoOUnVtcGkU3XX4Ws3GQfk= +go.opentelemetry.io/otel/sdk/log/logtest v0.19.0/go.mod h1:Lua81/3yM0wOmoHTokLj9y9ADeA02v1naRrVrkAZuKk= +go.opentelemetry.io/otel/sdk/metric v1.43.0 h1:S88dyqXjJkuBNLeMcVPRFXpRw2fuwdvfCGLEo89fDkw= +go.opentelemetry.io/otel/sdk/metric v1.43.0/go.mod h1:C/RJtwSEJ5hzTiUz5pXF1kILHStzb9zFlIEe85bhj6A= +go.opentelemetry.io/otel/trace v1.43.0 h1:BkNrHpup+4k4w+ZZ86CZoHHEkohws8AY+WTX09nk+3A= +go.opentelemetry.io/otel/trace v1.43.0/go.mod h1:/QJhyVBUUswCphDVxq+8mld+AvhXZLhe+8WVFxiFff0= +go.opentelemetry.io/proto/otlp v1.10.0 h1:IQRWgT5srOCYfiWnpqUYz9CVmbO8bFmKcwYxpuCSL2g= +go.opentelemetry.io/proto/otlp v1.10.0/go.mod h1:/CV4QoCR/S9yaPj8utp3lvQPoqMtxXdzn7ozvvozVqk= go.uber.org/goleak v1.3.0 h1:2K3zAYmnTNqV73imy9J1T3WC+gmCePx2hEGkimedGto= go.uber.org/goleak v1.3.0/go.mod h1:CoHD4mav9JJNrW/WLlf7HGZPjdw8EucARQHekz1X6bE= -golang.org/x/crypto v0.49.0 h1:+Ng2ULVvLHnJ/ZFEq4KdcDd/cfjrrjjNSXNzxg0Y4U4= -golang.org/x/crypto v0.49.0/go.mod h1:ErX4dUh2UM+CFYiXZRTcMpEcN8b/1gxEuv3nODoYtCA= -golang.org/x/net v0.52.0 h1:He/TN1l0e4mmR3QqHMT2Xab3Aj3L9qjbhRm78/6jrW0= -golang.org/x/net v0.52.0/go.mod h1:R1MAz7uMZxVMualyPXb+VaqGSa3LIaUqk0eEt3w36Sw= +go.yaml.in/yaml/v3 v3.0.4/go.mod h1:DhzuOOF2ATzADvBadXxruRBLzYTpT36CKvDb3+aBEFg= +golang.org/x/crypto v0.51.0 h1:IBPXwPfKxY7cWQZ38ZCIRPI50YLeevDLlLnyC5wRGTI= +golang.org/x/crypto v0.51.0/go.mod h1:8AdwkbraGNABw2kOX6YFPs3WM22XqI4EXEd8g+x7Oc8= +golang.org/x/mod v0.35.0 h1:Ww1D637e6Pg+Zb2KrWfHQUnH2dQRLBQyAtpr/haaJeM= +golang.org/x/mod v0.35.0/go.mod h1:+GwiRhIInF8wPm+4AoT6L0FA1QWAad3OMdTRx4tFYlU= +golang.org/x/net v0.55.0 h1:bcvxaJn3e1U6InsFWt1JUq1aSjnRxLzT2rtD2KfkDF8= +golang.org/x/net v0.55.0/go.mod h1:L5U2KuzuOe1lY7Z+aWVIKK6qEeJXnXV9yzGA+WCHJww= golang.org/x/oauth2 v0.36.0 h1:peZ/1z27fi9hUOFCAZaHyrpWG5lwe0RJEEEeH0ThlIs= golang.org/x/oauth2 v0.36.0/go.mod h1:YDBUJMTkDnJS+A4BP4eZBjCqtokkg1hODuPjwiGPO7Q= golang.org/x/sync v0.20.0 h1:e0PTpb7pjO8GAtTs2dQ6jYa5BWYlMuX047Dco/pItO4= golang.org/x/sync v0.20.0/go.mod h1:9xrNwdLfx4jkKbNva9FpL6vEN7evnE43NNNJQ2LF3+0= golang.org/x/sys v0.0.0-20220811171246-fbc7d0a398ab/go.mod h1:oPkhp1MJrh7nUepCBck5+mAzfO9JrbApNNgaTdGDITg= -golang.org/x/sys v0.42.0 h1:omrd2nAlyT5ESRdCLYdm3+fMfNFE/+Rf4bDIQImRJeo= -golang.org/x/sys v0.42.0/go.mod h1:4GL1E5IUh+htKOUEOaiffhrAeqysfVGipDYzABqnCmw= +golang.org/x/sys v0.45.0 h1:dO4czNzziLiiXplLQgBCEpCvXQ3dnkn0SdaZSYdQ+FY= +golang.org/x/sys v0.45.0/go.mod h1:4GL1E5IUh+htKOUEOaiffhrAeqysfVGipDYzABqnCmw= golang.org/x/text v0.3.3/go.mod h1:5Zoc/QRtKVWzQhOtBMvqHzDpF6irO9z98xDceosuGiQ= -golang.org/x/text v0.35.0 h1:JOVx6vVDFokkpaq1AEptVzLTpDe9KGpj5tR4/X+ybL8= -golang.org/x/text v0.35.0/go.mod h1:khi/HExzZJ2pGnjenulevKNX1W67CUy0AsXcNubPGCA= +golang.org/x/text v0.37.0 h1:Cqjiwd9eSg8e0QAkyCaQTNHFIIzWtidPahFWR83rTrc= +golang.org/x/text v0.37.0/go.mod h1:a5sjxXGs9hsn/AJVwuElvCAo9v8QYLzvavO5z2PiM38= golang.org/x/time v0.15.0 h1:bbrp8t3bGUeFOx08pvsMYRTCVSMk89u4tKbNOZbp88U= golang.org/x/time v0.15.0/go.mod h1:Y4YMaQmXwGQZoFaVFk4YpCt4FLQMYKZe9oeV/f4MSno= golang.org/x/tools v0.0.0-20180917221912-90fa682c2a6e/go.mod h1:n7NCudcB/nEzxVGmLbDWY5pfWTLqBcC2KZ6jyYvM4mQ= -golang.org/x/tools v0.42.0 h1:uNgphsn75Tdz5Ji2q36v/nsFSfR/9BRFvqhGBaJGd5k= -golang.org/x/tools v0.42.0/go.mod h1:Ma6lCIwGZvHK6XtgbswSoWroEkhugApmsXyrUmBhfr0= -gonum.org/v1/gonum v0.16.0 h1:5+ul4Swaf3ESvrOnidPp4GZbzf0mxVQpDCYUQE7OJfk= -gonum.org/v1/gonum v0.16.0/go.mod h1:fef3am4MQ93R2HHpKnLk4/Tbh/s0+wqD5nfa6Pnwy4E= -google.golang.org/api v0.272.0 h1:eLUQZGnAS3OHn31URRf9sAmRk3w2JjMx37d2k8AjJmA= -google.golang.org/api v0.272.0/go.mod h1:wKjowi5LNJc5qarNvDCvNQBn3rVK8nSy6jg2SwRwzIA= -google.golang.org/genai v1.40.0 h1:kYxyQSH+vsib8dvsgyLJzsVEIv5k3ZmHJyVqdvGncmc= -google.golang.org/genai v1.40.0/go.mod h1:A3kkl0nyBjyFlNjgxIwKq70julKbIxpSxqKO5gw/gmk= -google.golang.org/genproto v0.0.0-20260217215200-42d3e9bedb6d h1:vsOm753cOAMkt76efriTCDKjpCbK18XGHMJHo0JUKhc= -google.golang.org/genproto v0.0.0-20260217215200-42d3e9bedb6d/go.mod h1:0oz9d7g9QLSdv9/lgbIjowW1JoxMbxmBVNe8i6tORJI= -google.golang.org/genproto/googleapis/api v0.0.0-20260316180232-0b37fe3546d5 h1:CogIeEXn4qWYzzQU0QqvYBM8yDF9cFYzDq9ojSpv0Js= -google.golang.org/genproto/googleapis/api v0.0.0-20260316180232-0b37fe3546d5/go.mod h1:EIQZ5bFCfRQDV4MhRle7+OgjNtZ6P1PiZBgAKuxXu/Y= -google.golang.org/genproto/googleapis/rpc v0.0.0-20260316180232-0b37fe3546d5 h1:aJmi6DVGGIStN9Mobk/tZOOQUBbj0BPjZjjnOdoZKts= -google.golang.org/genproto/googleapis/rpc v0.0.0-20260316180232-0b37fe3546d5/go.mod h1:4Hqkh8ycfw05ld/3BWL7rJOSfebL2Q+DVDeRgYgxUU8= -google.golang.org/grpc v1.79.3 h1:sybAEdRIEtvcD68Gx7dmnwjZKlyfuc61Dyo9pGXXkKE= -google.golang.org/grpc v1.79.3/go.mod h1:KmT0Kjez+0dde/v2j9vzwoAScgEPx/Bw1CYChhHLrHQ= +golang.org/x/tools v0.44.0 h1:UP4ajHPIcuMjT1GqzDWRlalUEoY+uzoZKnhOjbIPD2c= +golang.org/x/tools v0.44.0/go.mod h1:KA0AfVErSdxRZIsOVipbv3rQhVXTnlU6UhKxHd1seDI= +gonum.org/v1/gonum v0.17.0 h1:VbpOemQlsSMrYmn7T2OUvQ4dqxQXU+ouZFQsZOx50z4= +gonum.org/v1/gonum v0.17.0/go.mod h1:El3tOrEuMpv2UdMrbNlKEh9vd86bmQ6vqIcDwxEOc1E= +google.golang.org/api v0.279.0 h1:hsx2M2OaRcaKtVYK6vXEUnQvdjnend7ZYES+lYaot74= +google.golang.org/api v0.279.0/go.mod h1:B9TqLBwJqVjp1mtt7WeoQwWRwvu/400y5lETOql+giQ= +google.golang.org/genai v1.57.0 h1:qTyG2ynz5dQy2jF4CvZdLHHVslhR0heMue+zM1a4GNM= +google.golang.org/genai v1.57.0/go.mod h1:A3kkl0nyBjyFlNjgxIwKq70julKbIxpSxqKO5gw/gmk= +google.golang.org/genproto v0.0.0-20260319201613-d00831a3d3e7 h1:XzmzkmB14QhVhgnawEVsOn6OFsnpyxNPRY9QV01dNB0= +google.golang.org/genproto v0.0.0-20260319201613-d00831a3d3e7/go.mod h1:L43LFes82YgSonw6iTXTxXUX1OlULt4AQtkik4ULL/I= +google.golang.org/genproto/googleapis/api v0.0.0-20260427160629-7cedc36a6bc4 h1:yOzSCGPx+cp5VO7IxvZ9SBFF7j1tZVcNtlHR2iYKtVo= +google.golang.org/genproto/googleapis/api v0.0.0-20260427160629-7cedc36a6bc4/go.mod h1:Q9HWtNeE7tM9npdIsEvqXj1QJIvVoeAV3rtXtS715Cw= +google.golang.org/genproto/googleapis/rpc v0.0.0-20260511170946-3700d4141b60 h1:seT2EwLWM78plQ7wcDfuWBc/4FAEAXDDiaSol4ku4qo= +google.golang.org/genproto/googleapis/rpc v0.0.0-20260511170946-3700d4141b60/go.mod h1:4Hqkh8ycfw05ld/3BWL7rJOSfebL2Q+DVDeRgYgxUU8= +google.golang.org/grpc v1.81.0 h1:W3G9N3KQf3BU+YuCtGKJk0CmxQNbAISICD/9AORxLIw= +google.golang.org/grpc v1.81.0/go.mod h1:xGH9GfzOyMTGIOXBJmXt+BX/V0kcdQbdcuwQ/zNw42I= google.golang.org/protobuf v1.36.11 h1:fV6ZwhNocDyBLK0dj+fg8ektcVegBBuEolpbTQyBNVE= google.golang.org/protobuf v1.36.11/go.mod h1:HTf+CrKn2C3g5S8VImy6tdcUvCska2kB7j23XfzDpco= gopkg.in/check.v1 v0.0.0-20161208181325-20d25e280405/go.mod h1:Co6ibVJAznAaIkqp8huTwlJQCZ016jof/cbN4VW5Yz0= diff --git a/internal/agent/remoteagent/a2a_config.go b/internal/agent/remoteagent/a2a_config.go index ef5dd92ec..84b8e55c1 100644 --- a/internal/agent/remoteagent/a2a_config.go +++ b/internal/agent/remoteagent/a2a_config.go @@ -16,16 +16,31 @@ package remoteagent import ( "context" - "encoding/json" "fmt" - "os" - "strings" + "iter" - "github.com/a2aproject/a2a-go/a2a" - "github.com/a2aproject/a2a-go/a2aclient" - "github.com/a2aproject/a2a-go/a2aclient/agentcard" + "github.com/a2aproject/a2a-go/v2/a2a" + "github.com/a2aproject/a2a-go/v2/a2aclient" ) +// A2AClient abstracts a2a-go client so that we can use different SDK versions. +type A2AClient interface { + // SendMessage sends a message to the remote agent and returns the result. + SendMessage(ctx context.Context, req *a2a.SendMessageRequest) (a2a.SendMessageResult, error) + // SendStreamingMessage sends a message to the remote agent and returns a stream of events. + SendStreamingMessage(ctx context.Context, req *a2a.SendMessageRequest) iter.Seq2[a2a.Event, error] + // CancelTask cancels a task on the remote agent. + CancelTask(ctx context.Context, req *a2a.CancelTaskRequest) (*a2a.Task, error) + // Destroy is called in the end of agent invocation. + Destroy() error +} + +// A2AClientProvider creates an [A2AClient]. +type A2AClientProvider interface { + // CreateClient creates an [A2AClient]. + CreateClient(context.Context, *a2a.AgentCard) (A2AClient, error) +} + // RemoteAgentState holds the internal state of a remote agent. type RemoteAgentState struct { // A2A holds the A2A configuration if remote agent is an A2A agent. @@ -34,25 +49,23 @@ type RemoteAgentState struct { // A2AServerConfig is used to describe and configure a remote agent. type A2AServerConfig struct { - // AgentCardSource can be either an http(s) URL or a local file path. If a2a.AgentCard - // is not provided, the source is used to resolve the card during the first agent invocation. - AgentCard *a2a.AgentCard - AgentCardSource string - // CardResolveOptions can be used to provide a set of agencard.Resolver configurations. - CardResolveOptions []agentcard.ResolveOption - // ClientFactory can be used to provide a set of a2aclient.Client configurations. - ClientFactory *a2aclient.Factory + // AgentCard is a static agent card. + AgentCard *a2a.AgentCard + // AgentCardProvider resolves an agent card lazily. + AgentCardProvider func(ctx context.Context) (*a2a.AgentCard, error) + // ClientProvider is used to create an [A2AClient] implementation. + ClientProvider A2AClientProvider } -func CreateA2AClient(ctx context.Context, cfg *A2AServerConfig) (*a2a.AgentCard, *a2aclient.Client, error) { - card, err := resolveAgentCard(ctx, cfg) +func CreateA2AClient(ctx context.Context, cfg *A2AServerConfig) (*a2a.AgentCard, A2AClient, error) { + card, err := ResolveAgentCard(ctx, cfg) if err != nil { return nil, nil, fmt.Errorf("agent card resolution failed: %w", err) } - var client *a2aclient.Client - if cfg.ClientFactory != nil { - client, err = cfg.ClientFactory.CreateFromCard(ctx, card) + var client A2AClient + if cfg.ClientProvider != nil { + client, err = cfg.ClientProvider.CreateClient(ctx, card) } else { client, err = a2aclient.NewFromCard(ctx, card) } @@ -62,27 +75,12 @@ func CreateA2AClient(ctx context.Context, cfg *A2AServerConfig) (*a2a.AgentCard, return card, client, nil } -func resolveAgentCard(ctx context.Context, cfg *A2AServerConfig) (*a2a.AgentCard, error) { +func ResolveAgentCard(ctx context.Context, cfg *A2AServerConfig) (*a2a.AgentCard, error) { if cfg.AgentCard != nil { return cfg.AgentCard, nil } - - if strings.HasPrefix(cfg.AgentCardSource, "http://") || strings.HasPrefix(cfg.AgentCardSource, "https://") { - card, err := agentcard.DefaultResolver.Resolve(ctx, cfg.AgentCardSource, cfg.CardResolveOptions...) - if err != nil { - return nil, fmt.Errorf("failed to fetch an agent card: %w", err) - } - return card, nil - } - - fileBytes, err := os.ReadFile(cfg.AgentCardSource) - if err != nil { - return nil, fmt.Errorf("failed to read agent card from %q: %w", cfg.AgentCardSource, err) - } - - var card a2a.AgentCard - if err := json.Unmarshal(fileBytes, &card); err != nil { - return nil, fmt.Errorf("failed to unmarshal an agent card: %w", err) + if cfg.AgentCardProvider != nil { + return cfg.AgentCardProvider(ctx) } - return &card, nil + return nil, fmt.Errorf("either AgentCard or AgentCardProvider must be set") } diff --git a/internal/agent/runconfig/run_config.go b/internal/agent/runconfig/run_config.go index ecf748a5b..325451d57 100644 --- a/internal/agent/runconfig/run_config.go +++ b/internal/agent/runconfig/run_config.go @@ -14,7 +14,11 @@ package runconfig -import "context" +import ( + "context" + + "google.golang.org/adk/agent" +) type StreamingMode string @@ -26,6 +30,7 @@ const ( type RunConfig struct { StreamingMode StreamingMode + Live *agent.LiveRunConfig } func ToContext(ctx context.Context, cfg *RunConfig) context.Context { diff --git a/internal/configurable/conformance/functions.go b/internal/configurable/conformance/functions.go index a89eceff6..22445a57e 100644 --- a/internal/configurable/conformance/functions.go +++ b/internal/configurable/conformance/functions.go @@ -21,6 +21,7 @@ import ( "math" "regexp" + "google.golang.org/adk/agent" "google.golang.org/adk/internal/configurable" "google.golang.org/adk/tool" "google.golang.org/adk/tool/functiontool" @@ -32,11 +33,11 @@ type ValidateEmailArgs struct { var emailRegex = regexp.MustCompile(`^[a-zA-Z0-9._%+-]+@[a-zA-Z0-9.-]+\.[a-zA-Z]{2,}$`) -func validateEmail(ctx tool.Context, args ValidateEmailArgs) (bool, error) { +func validateEmail(ctx agent.ToolContext, args ValidateEmailArgs) (bool, error) { return emailRegex.MatchString(args.Email), nil } -func getUserID(ctx tool.Context, args ValidateEmailArgs) (int, error) { +func getUserID(ctx agent.ToolContext, args ValidateEmailArgs) (int, error) { valid, err := validateEmail(ctx, args) if err != nil { return 0, err @@ -58,15 +59,17 @@ func getUserID(ctx tool.Context, args ValidateEmailArgs) (int, error) { return int(result % 10000), nil } -func createBooking(ctx tool.Context, args ValidateEmailArgs) (map[string]any, error) { - userID, err := getUserID(ctx, args) - if err != nil { - return nil, err - } +type CreateBookingArgs struct { + UserID int `json:"user_id"` + IsConfirmed bool `json:"is_confirmed"` + Details string `json:"details"` +} + +func createBooking(ctx agent.ToolContext, args CreateBookingArgs) (map[string]any, error) { return map[string]any{ - "user_id": userID, - "is_confirmed": true, - "details": "Booking created for user " + args.Email, + "user_id": args.UserID, + "is_confirmed": args.IsConfirmed, + "details": args.Details, "user_id_type": "int", "is_confirmed_type": "bool", "details_type": "string", @@ -92,7 +95,7 @@ type SearchFlightsArgs struct { Preferences *FlightPreferences `json:"preferences"` } -func searchFlights(ctx tool.Context, args SearchFlightsArgs) (map[string]any, error) { +func searchFlights(ctx agent.ToolContext, args SearchFlightsArgs) (map[string]any, error) { if args.Preferences == nil { args.Preferences = &FlightPreferences{ CabinClass: "Economy", @@ -150,7 +153,7 @@ type CalculateTripCostArgs struct { BaggageCount *int `json:"baggage_count"` } -func calculateTripCost(ctx tool.Context, args CalculateTripCostArgs) (map[string]any, error) { +func calculateTripCost(ctx agent.ToolContext, args CalculateTripCostArgs) (map[string]any, error) { // Handle Python's default num_passengers=1 logic // In Go, if the caller passes 0, we should ensure at least 1 // or handle it based on your specific business logic. @@ -198,7 +201,7 @@ type reimburseArgs struct { Amount float64 `json:"amount"` } -func reimburse(ctx tool.Context, args reimburseArgs) (map[string]any, error) { +func reimburse(ctx agent.ToolContext, args reimburseArgs) (map[string]any, error) { return map[string]any{ "status": "ok", }, nil @@ -209,7 +212,7 @@ type askForApprovalArgs struct { Amount float64 `json:"amount"` } -func askForApproval(ctx tool.Context, args askForApprovalArgs) (map[string]any, error) { +func askForApproval(ctx agent.ToolContext, args askForApprovalArgs) (map[string]any, error) { return map[string]any{ "status": "pending", "amount": args.Amount, @@ -243,8 +246,7 @@ Args: Returns: A dictionary containing the booking information and the types of the - received arguments. -`, + received arguments.`, }, createBooking) if err != nil { return fmt.Errorf("error creating create booking tool: %w", err) diff --git a/internal/configurable/conformance/recordplugin/invocation_record_state.go b/internal/configurable/conformance/recordplugin/invocation_record_state.go new file mode 100644 index 000000000..ed776e624 --- /dev/null +++ b/internal/configurable/conformance/recordplugin/invocation_record_state.go @@ -0,0 +1,114 @@ +// Copyright 2026 Google LLC +// +// Licensed under the Apache License, Version 2.0 (the "License"); +// you may not use this file except in compliance with the License. +// You may obtain a copy of the License at +// +// http://www.apache.org/licenses/LICENSE-2.0 +// +// Unless required by applicable law or agreed to in writing, software +// distributed under the License is distributed on an "AS IS" BASIS, +// WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. +// See the License for the specific language governing permissions and +// limitations under the License. + +package recordplugin + +import ( + "sync" + + "google.golang.org/genai" + + "google.golang.org/adk/internal/configurable/conformance/replayplugin/recording" + "google.golang.org/adk/model" +) + +type invocationRecordState struct { + mu sync.Mutex + caseDir string + userMessageIndex int + streamingMode string + + // Completed recordings for the current invocation turn + recordings []recording.Recording + + // Active/pending LLM request mappings + pendingLLMRequest *model.LLMRequest + pendingLLMResponses []*model.LLMResponse + + // Active/pending Tool call mappings keyed by function_call_id + pendingToolCalls map[string]*genai.FunctionCall +} + +func newInvocationRecordState(caseDir string, msgIndex int, streamingMode string) *invocationRecordState { + return &invocationRecordState{ + caseDir: caseDir, + userMessageIndex: msgIndex, + streamingMode: streamingMode, + pendingToolCalls: make(map[string]*genai.FunctionCall), + } +} + +func (s *invocationRecordState) StartLLMRecording(req *model.LLMRequest) { + s.mu.Lock() + defer s.mu.Unlock() + + s.pendingLLMRequest = req + s.pendingLLMResponses = nil +} + +func (s *invocationRecordState) AppendLLMResponse(resp *model.LLMResponse) { + s.mu.Lock() + defer s.mu.Unlock() + + s.pendingLLMResponses = append(s.pendingLLMResponses, resp) +} + +func (s *invocationRecordState) CompleteLLMRecording(agentName string) { + s.mu.Lock() + defer s.mu.Unlock() + + if s.pendingLLMRequest == nil { + return + } + + s.recordings = append(s.recordings, recording.Recording{ + UserMessageIndex: s.userMessageIndex, + AgentName: agentName, + LLMRecording: &recording.LLMRecording{ + LLMRequest: s.pendingLLMRequest, + LLMResponses: s.pendingLLMResponses, + }, + }) + + s.pendingLLMRequest = nil + s.pendingLLMResponses = nil +} + +func (s *invocationRecordState) StartToolRecording(id string, fc *genai.FunctionCall) { + s.mu.Lock() + defer s.mu.Unlock() + + s.pendingToolCalls[id] = fc +} + +func (s *invocationRecordState) CompleteToolRecording(id, agentName string, resp *genai.FunctionResponse) { + s.mu.Lock() + defer s.mu.Unlock() + + fc, ok := s.pendingToolCalls[id] + if !ok { + return + } + + s.recordings = append(s.recordings, recording.Recording{ + UserMessageIndex: s.userMessageIndex, + AgentName: agentName, + ToolRecording: &recording.ToolRecording{ + ToolCall: fc, + ToolResponse: resp, + }, + }) + + delete(s.pendingToolCalls, id) +} diff --git a/internal/configurable/conformance/recordplugin/record_plugin.go b/internal/configurable/conformance/recordplugin/record_plugin.go new file mode 100644 index 000000000..ee73d4697 --- /dev/null +++ b/internal/configurable/conformance/recordplugin/record_plugin.go @@ -0,0 +1,625 @@ +// Copyright 2026 Google LLC +// +// Licensed under the Apache License, Version 2.0 (the "License"); +// you may not use this file except in compliance with the License. +// You may obtain a copy of the License at +// +// http://www.apache.org/licenses/LICENSE-2.0 +// +// Unless required by applicable law or agreed to in writing, software +// distributed under the License is distributed on an "AS IS" BASIS, +// WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. +// See the License for the specific language governing permissions and +// limitations under the License. + +package recordplugin + +import ( + "encoding/base64" + "fmt" + "os" + "path/filepath" + "strings" + "sync" + + "google.golang.org/genai" + "gopkg.in/yaml.v3" + + "google.golang.org/adk/agent" + "google.golang.org/adk/internal/configurable/conformance/replayplugin/recording" + "google.golang.org/adk/model" + "google.golang.org/adk/plugin" + "google.golang.org/adk/session" + "google.golang.org/adk/tool" +) + +type recordPlugin struct { + mu sync.Mutex + invocationStates map[string]*invocationRecordState + allowedBaseDir string +} + +// New creates a new conformance record plugin. +func New(allowedBaseDir string) (*plugin.Plugin, error) { + p := &recordPlugin{ + invocationStates: make(map[string]*invocationRecordState), + allowedBaseDir: allowedBaseDir, + } + return plugin.New(plugin.Config{ + Name: "record_plugin", + BeforeRunCallback: p.beforeRun, + AfterRunCallback: p.afterRun, + BeforeModelCallback: p.beforeModel, + AfterModelCallback: p.afterModel, + BeforeToolCallback: p.beforeTool, + AfterToolCallback: p.afterTool, + }) +} + +// MustNew is like New but panics on error. +func MustNew(allowedBaseDir string) *plugin.Plugin { + p, err := New(allowedBaseDir) + if err != nil { + panic(err) + } + return p +} + +func (p *recordPlugin) beforeRun(ctx agent.InvocationContext) (*genai.Content, error) { + if ctx.Session() == nil { + return nil, nil + } + + on, err := p.isRecordModeOn(ctx.Session().State()) + if err != nil { + return nil, err + } + if !on { + return nil, nil + } + + // Create fresh record state for this invocation + _, err = p.createInvocationState(ctx) + return nil, err +} + +func (p *recordPlugin) beforeModel(ctx agent.CallbackContext, req *model.LLMRequest) (*model.LLMResponse, error) { + on, err := p.isRecordModeOn(ctx.State()) + if err != nil || !on { + return nil, nil + } + + state, err := p.getInvocationState(ctx.InvocationID()) + if err != nil { + return nil, err + } + + // Sanitize and clone the request to prevent YAML serialization panics caused by HTTPOptions.ExtrasRequestProvider + sanitizedReq := sanitizeLLMRequest(req) + + state.StartLLMRecording(sanitizedReq) + return nil, nil +} + +func (p *recordPlugin) afterModel(ctx agent.CallbackContext, resp *model.LLMResponse, err error) (*model.LLMResponse, error) { + on, recordErr := p.isRecordModeOn(ctx.State()) + if recordErr != nil || !on { + return nil, nil + } + + if err != nil || resp == nil { + return nil, nil + } + + state, stateErr := p.getInvocationState(ctx.InvocationID()) + if stateErr != nil { + return nil, stateErr + } + + state.AppendLLMResponse(resp) + + // If it is a final non-partial response, complete the recording + if !resp.Partial { + state.CompleteLLMRecording(ctx.AgentName()) + } + + return nil, nil +} + +func (p *recordPlugin) beforeTool(ctx agent.ToolContext, t tool.Tool, args map[string]any) (map[string]any, error) { + on, err := p.isRecordModeOn(ctx.State()) + if err != nil || !on { + return nil, nil + } + + state, err := p.getInvocationState(ctx.InvocationID()) + if err != nil { + return nil, err + } + + fc := &genai.FunctionCall{ + ID: ctx.FunctionCallID(), + Name: t.Name(), + Args: args, + } + + state.StartToolRecording(ctx.FunctionCallID(), fc) + return nil, nil +} + +func (p *recordPlugin) afterTool(ctx agent.ToolContext, t tool.Tool, args, result map[string]any, err error) (map[string]any, error) { + on, recordErr := p.isRecordModeOn(ctx.State()) + if recordErr != nil || !on { + return nil, nil + } + + state, stateErr := p.getInvocationState(ctx.InvocationID()) + if stateErr != nil { + return nil, stateErr + } + + resp := &genai.FunctionResponse{ + ID: ctx.FunctionCallID(), + Name: t.Name(), + Response: result, + } + + state.CompleteToolRecording(ctx.FunctionCallID(), ctx.AgentName(), resp) + return nil, nil +} + +func (p *recordPlugin) afterRun(ctx agent.InvocationContext) { + if ctx.Session() == nil { + return + } + + on, err := p.isRecordModeOn(ctx.Session().State()) + if err != nil || !on { + return + } + + state, err := p.getInvocationState(ctx.InvocationID()) + if err != nil { + return + } + + defer func() { + p.mu.Lock() + delete(p.invocationStates, ctx.InvocationID()) + p.mu.Unlock() + }() + + // Serialize and save + p.saveRecordings(state) +} + +func (p *recordPlugin) parseRecordConfig(sessionState session.State) (string, int, string, error) { + if sessionState == nil { + return "", 0, "", nil + } + + configVal, err := sessionState.Get("_adk_recordings_config") + if err != nil { + return "", 0, "", nil + } + + config, ok := configVal.(map[string]any) + if !ok { + return "", 0, "", nil + } + + caseDir, ok := config["dir"].(string) + if !ok || caseDir == "" { + return "", 0, "", nil + } + + basePath, err := filepath.Abs(p.allowedBaseDir) + if err != nil { + return "", 0, "", fmt.Errorf("invalid path format: %v", err) + } + requestedAbsPath, err := filepath.Abs(caseDir) + if err != nil { + return "", 0, "", fmt.Errorf("invalid path format: %v", err) + } + rel, err := filepath.Rel(basePath, requestedAbsPath) + if err != nil { + return "", 0, "", fmt.Errorf("invalid path format: %v", err) + } + if strings.HasPrefix(rel, "..") || filepath.IsAbs(rel) { + return "", 0, "", fmt.Errorf("record config error: 'dir' is not within the allowed base directory") + } + + msgIndexVal, ok := config["user_message_index"] + if !ok || msgIndexVal == nil { + return "", 0, "", nil + } + + var msgIndex int + switch v := msgIndexVal.(type) { + case int: + msgIndex = v + case float64: + msgIndex = int(v) + default: + return "", 0, "", fmt.Errorf("record config 'user_message_index' is not a number") + } + + streamingMode, _ := config["streaming_mode"].(string) + if streamingMode == "" { + streamingMode = "none" + } + + return caseDir, msgIndex, streamingMode, nil +} + +func (p *recordPlugin) isRecordModeOn(sessionState session.State) (bool, error) { + caseDir, _, _, err := p.parseRecordConfig(sessionState) + if err != nil { + return false, err + } + return caseDir != "", nil +} + +func (p *recordPlugin) getInvocationState(id string) (*invocationRecordState, error) { + p.mu.Lock() + defer p.mu.Unlock() + + state, ok := p.invocationStates[id] + if !ok { + return nil, fmt.Errorf("record state not initialized for invocation: %s", id) + } + return state, nil +} + +func (p *recordPlugin) createInvocationState(ctx agent.InvocationContext) (*invocationRecordState, error) { + caseDir, msgIndex, streamingMode, err := p.parseRecordConfig(ctx.Session().State()) + if err != nil { + return nil, err + } + if caseDir == "" { + return nil, fmt.Errorf("record state not configured") + } + + state := newInvocationRecordState(caseDir, msgIndex, streamingMode) + + p.mu.Lock() + p.invocationStates[ctx.InvocationID()] = state + p.mu.Unlock() + + return state, nil +} + +func (p *recordPlugin) saveRecordings(state *invocationRecordState) { + var filename string + if state.streamingMode == "sse" { + filename = "generated-recordings-sse.yaml" + } else { + filename = "generated-recordings.yaml" + } + + filePath := filepath.Join(state.caseDir, filename) + + // Load existing recordings if the file exists + var existingRecordings recording.Recordings + if _, err := os.Stat(filePath); err == nil { + data, err := os.ReadFile(filePath) + if err == nil { + _ = yaml.Unmarshal(data, &existingRecordings) + } + } + + // Filter out any existing recordings for the CURRENT user_message_index + var filtered []recording.Recording + for _, r := range existingRecordings.Recordings { + if r.UserMessageIndex != state.userMessageIndex { + filtered = append(filtered, r) + } + } + + // Append newly completed recordings in this turn + state.mu.Lock() + filtered = append(filtered, state.recordings...) + state.mu.Unlock() + + outputRecordings := recording.Recordings{ + Recordings: filtered, + } + + // Write YAML back to disk using AST traversal to prune nulls/defaults + var node yaml.Node + if err := node.Encode(&outputRecordings); err != nil { + return + } + + cleanYAMLNode(&node, false) + + out, err := yaml.Marshal(&node) + if err != nil { + return + } + + _ = os.WriteFile(filePath, out, 0o644) +} + +func shouldPruneYAMLField(key string, value *yaml.Node, inToolArgs bool) bool { + if key == "usermessageindex" || key == "user_message_index" { + return false + } + if !inToolArgs { + if key == "property_order" || key == "response_json_schema" { + return true + } + } + if value.Kind == yaml.ScalarNode { + if value.Tag == "!!null" || value.Value == "null" || value.Value == "~" { + return true + } + if !inToolArgs { + if (value.Tag == "!!int" || value.Tag == "") && value.Value == "0" { + return true + } + if (value.Tag == "!!bool" || value.Tag == "") && value.Value == "false" { + return true + } + if value.Value == "" { + return true + } + } + } + if value.Kind == yaml.SequenceNode && len(value.Content) == 0 { + return true + } + if value.Kind == yaml.MappingNode && len(value.Content) == 0 { + if inToolArgs { + return false + } + switch key { + case "google_search", "code_execution", "retrieval": + return false + default: + return true + } + } + return false +} + +var snakeCaseFieldMap = make(map[string]string) + +func init() { + snakeCaseKeys := []string{ + // Top-level and Recording structs + "user_message_index", "agent_name", "llm_recording", "llm_request", "llm_responses", + "tool_recording", "tool_call", "tool_response", + + // LLMResponse fields + "citation_metadata", "grounding_metadata", "usage_metadata", "custom_metadata", + "logprobs_result", "input_transcription", "output_transcription", "model_version", + "session_resumption_handle", "error_code", "error_message", "finish_reason", "avg_logprobs", + + // GenerateContentConfig fields + "http_options", "system_instruction", "candidate_count", "max_output_tokens", + "stop_sequences", "response_logprobs", "presence_penalty", "frequency_penalty", + "response_mime_type", "response_schema", "response_json_schema", "routing_config", + "model_selection_config", "safety_settings", "tool_config", "cached_content", + "response_modalities", "media_resolution", "speech_config", "audio_timestamp", + "thinking_config", "image_config", "enable_enhanced_civic_answers", "model_armor_config", + "service_tier", + + // Tool fields + "google_search", "blocking_confidence", "function_declarations", "parameters_json_schema", + "additional_properties", "property_order", + + // Part / Content fields + "part_metadata", "video_metadata", "code_execution_result", "executable_code", + "file_data", "function_call", "function_call_id", "function_response", "inline_data", + "thought_signature", + + // GroundingMetadata fields + "grounding_chunks", "grounding_supports", "retrieval_metadata", "search_entry_point", + "web_search_queries", + + // GroundingChunk / Support / Segment / SearchEntryPoint fields + "grounding_chunk_indices", "confidence_score", "start_index", "end_index", + "rendered_content", "sdk_active_queries", + + // Token / usage details + "google_maps_widget_context_token", "candidates_token_count", "prompt_token_count", + "prompt_tokens_details", "token_count", "thoughts_token_count", "tool_use_prompt_token_count", + "tool_use_prompt_tokens_details", "total_token_count", "traffic_type", + } + + for _, key := range snakeCaseKeys { + cleaned := strings.ReplaceAll(key, "_", "") + snakeCaseFieldMap[cleaned] = key + } + + // Manual edge cases + snakeCaseFieldMap["propertyordering"] = "property_order" +} + +// findKeyNode searches a MappingNode's children to locate a value node associated +// with the specified key string. Returns nil if the node is not a mapping or if the key is not found. +func findKeyNode(node *yaml.Node, key string) *yaml.Node { + if node.Kind != yaml.MappingNode { + return nil + } + for i := 0; i < len(node.Content); i += 2 { + if node.Content[i].Value == key { + return node.Content[i+1] + } + } + return nil +} + +// cleanYAMLNode is the a post-processor for the recorded YAML representation. +// It performs a deep recursive AST traversal to: +// +// Normalize Go-serialized lowercase keys back to their strict canonical snake_case. +// Auto-extract and simplify standard system_instruction structures to single scalar strings. +// Coalesce raw integer-sequence thought_signatures back to compact Base64 strings. +// Prune all default empty fields, empty sequences, and null properties outside user tool arguments. +func cleanYAMLNode(node *yaml.Node, inToolArgs bool) { + switch node.Kind { + case yaml.DocumentNode, yaml.SequenceNode: + for _, child := range node.Content { + cleanYAMLNode(child, inToolArgs) + } + case yaml.MappingNode: + // If this mapping represents an OpenAPI function declaration, recursively inject parameter schema titles. + // This ensures the generated schemas exactly match Pydantic reference standards. + var toolName string + if nameNode := findKeyNode(node, "name"); nameNode != nil && nameNode.Kind == yaml.ScalarNode { + toolName = nameNode.Value + } + var paramSchemaNode *yaml.Node + for _, key := range []string{"parameters", "parameters_json_schema", "parametersjsonschema"} { + if p := findKeyNode(node, key); p != nil && p.Kind == yaml.MappingNode { + paramSchemaNode = p + break + } + } + if toolName != "" && paramSchemaNode != nil { + injectSchemaTitles(paramSchemaNode, toolName) + } + + var newContent []*yaml.Node + for i := 0; i < len(node.Content); i += 2 { + k := node.Content[i] + v := node.Content[i+1] + + // Translate lowercase serialized struct fields back to canonical snake_case + if mappedKey, exists := snakeCaseFieldMap[k.Value]; exists { + k.Value = mappedKey + } + + // Set insideToolArgs context flag so we avoid pruning user data and parameter payloads + nextInToolArgs := inToolArgs + if k.Value == "args" || k.Value == "response" || k.Value == "tool_call" || k.Value == "tool_response" { + nextInToolArgs = true + } + + // Extract structured system instructions (Content with Parts) into simple string scalars + if k.Value == "system_instruction" && v.Kind == yaml.MappingNode { + if parts := findKeyNode(v, "parts"); parts != nil && parts.Kind == yaml.SequenceNode && len(parts.Content) > 0 { + part := parts.Content[0] + if text := findKeyNode(part, "text"); text != nil && text.Kind == yaml.ScalarNode { + v.Kind = yaml.ScalarNode + v.Tag = "!!str" + v.Style = text.Style + v.Value = text.Value + v.Content = nil + } + } + } + + // Coalesce raw integer sequences representing byte signatures back to standard Base64 string representation + if k.Value == "thought_signature" && v.Kind == yaml.SequenceNode && len(v.Content) > 0 { + var bytes []byte + for _, child := range v.Content { + if child.Kind == yaml.ScalarNode && (child.Tag == "!!int" || child.Tag == "") { + var b int + _, err := fmt.Sscan(child.Value, &b) + if err == nil { + bytes = append(bytes, byte(b)) + } + } + } + v.Kind = yaml.ScalarNode + v.Style = yaml.DoubleQuotedStyle + v.Tag = "!!str" + v.Value = base64.StdEncoding.EncodeToString(bytes) + v.Content = nil + } + + cleanYAMLNode(v, nextInToolArgs) + + // Drop empty, null, or redundant fields outside tool parameters + if shouldPruneYAMLField(k.Value, v, nextInToolArgs) { + continue + } + + newContent = append(newContent, k, v) + } + node.Content = newContent + } +} + +// injectSchemaTitles recursively verifies and inserts correct parameter schema titles +// for OpenAPI declarations, such as mapping "validate_email" to "validate_emailParams" +// and camelCasing individual property titles (e.g. "email_address" -> "Email Address"). +func injectSchemaTitles(schemaNode *yaml.Node, toolName string) { + if schemaNode.Kind != yaml.MappingNode { + return + } + + // Check or inject schema top-level title + titleNode := findKeyNode(schemaNode, "title") + if titleNode != nil { + if titleNode.Value == "" || titleNode.Tag == "!!null" || titleNode.Value == "null" { + titleNode.Kind = yaml.ScalarNode + titleNode.Tag = "!!str" + titleNode.Style = yaml.DoubleQuotedStyle + titleNode.Value = toolName + "Params" + titleNode.Content = nil + } + } else { + titleKey := &yaml.Node{Kind: yaml.ScalarNode, Tag: "!!str", Value: "title"} + titleValue := &yaml.Node{Kind: yaml.ScalarNode, Tag: "!!str", Value: toolName + "Params", Style: yaml.DoubleQuotedStyle} + schemaNode.Content = append([]*yaml.Node{titleKey, titleValue}, schemaNode.Content...) + } + + // Recursively check and inject titles for individual properties + if propertiesNode := findKeyNode(schemaNode, "properties"); propertiesNode != nil && propertiesNode.Kind == yaml.MappingNode { + for i := 0; i < len(propertiesNode.Content); i += 2 { + propNameNode := propertiesNode.Content[i] + propDefNode := propertiesNode.Content[i+1] + if propDefNode.Kind == yaml.MappingNode { + propTitle := findKeyNode(propDefNode, "title") + if propTitle != nil { + if propTitle.Value == "" || propTitle.Tag == "!!null" || propTitle.Value == "null" { + propTitle.Kind = yaml.ScalarNode + propTitle.Tag = "!!str" + propTitle.Style = yaml.DoubleQuotedStyle + propTitle.Value = toCamelCaseTitle(propNameNode.Value) + propTitle.Content = nil + } + } else { + titleKey := &yaml.Node{Kind: yaml.ScalarNode, Tag: "!!str", Value: "title"} + titleValue := &yaml.Node{Kind: yaml.ScalarNode, Tag: "!!str", Value: toCamelCaseTitle(propNameNode.Value), Style: yaml.DoubleQuotedStyle} + propDefNode.Content = append([]*yaml.Node{titleKey, titleValue}, propDefNode.Content...) + } + } + } + } +} + +// toCamelCaseTitle translates a snake_case property key (e.g. "phone_number") +// into a human-readable Title Case string (e.g. "Phone Number") for schema metadata. +func toCamelCaseTitle(s string) string { + parts := strings.Split(s, "_") + for i, part := range parts { + if len(part) > 0 { + parts[i] = strings.ToUpper(part[:1]) + part[1:] + } + } + return strings.Join(parts, " ") +} + +// sanitizeLLMRequest creates a shallow copy of the LLMRequest and strips HTTPOptions +// and internal Tool registries to prevent YAML serialization panics and leaked internal state. +func sanitizeLLMRequest(req *model.LLMRequest) *model.LLMRequest { + if req == nil { + return nil + } + + reqCopy := *req + reqCopy.Tools = nil + + if req.Config != nil { + configCopy := *req.Config + configCopy.HTTPOptions = nil + reqCopy.Config = &configCopy + } + + return &reqCopy +} diff --git a/internal/configurable/conformance/recordplugin/record_plugin_test.go b/internal/configurable/conformance/recordplugin/record_plugin_test.go new file mode 100644 index 000000000..60f06b5f5 --- /dev/null +++ b/internal/configurable/conformance/recordplugin/record_plugin_test.go @@ -0,0 +1,525 @@ +// Copyright 2026 Google LLC +// +// Licensed under the Apache License, Version 2.0 (the "License"); +// you may not use this file except in compliance with the License. +// You may obtain a copy of the License at +// +// http://www.apache.org/licenses/LICENSE-2.0 +// +// Unless required by applicable law or agreed to in writing, software +// distributed under the License is distributed on an "AS IS" BASIS, +// WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. +// See the License for the specific language governing permissions and +// limitations under the License. + +package recordplugin_test + +import ( + "context" + "iter" + "os" + "path/filepath" + "strings" + "testing" + "time" + + "google.golang.org/genai" + "gopkg.in/yaml.v3" + + "google.golang.org/adk/agent" + "google.golang.org/adk/internal/configurable/conformance/recordplugin" + "google.golang.org/adk/internal/configurable/conformance/replayplugin/recording" + "google.golang.org/adk/memory" + "google.golang.org/adk/model" + "google.golang.org/adk/plugin" + "google.golang.org/adk/session" + "google.golang.org/adk/tool/toolconfirmation" +) + +func TestRecordPlugin(t *testing.T) { + setup := func(t *testing.T, baseDir string) (*plugin.Plugin, *MockSession, *MockState) { + p := recordplugin.MustNew(baseDir) + sessionState := make(map[string]any) + mockState := &MockState{data: sessionState} + mockSession := &MockSession{state: mockState} + return p, mockSession, mockState + } + + t.Run("RecordLlmAndToolCallsAndSaveYaml", func(t *testing.T) { + tempDir := t.TempDir() + p, mockSession, _ := setup(t, tempDir) + + err := mockSession.State().Set("_adk_recordings_config", map[string]any{ + "dir": tempDir, + "user_message_index": 0, + }) + if err != nil { + t.Fatalf("unexpected error: %v", err) + } + + invContext := &MockInvocationContext{ + session: mockSession, + invocationID: "test-invocation", + } + + // 1. beforeRun + _, err = p.BeforeRunCallback()(invContext) + if err != nil { + t.Fatalf("beforeRun failed: %v", err) + } + + // 2. beforeModel (LLM request) + cbContext := &MockCallbackContext{ + state: mockSession.State(), + invocationID: "test-invocation", + agentName: "test_agent", + } + req := &model.LLMRequest{ + Model: "gemini-2.0-flash", + Contents: []*genai.Content{ + { + Role: "user", + Parts: []*genai.Part{{Text: "Hello agent"}}, + }, + }, + Config: &genai.GenerateContentConfig{ + HTTPOptions: &genai.HTTPOptions{ + ExtrasRequestProvider: func(body map[string]any) map[string]any { + return body + }, + }, + }, + Tools: map[string]any{ + "test_tool": "dummy_val", + }, + } + _, err = p.BeforeModelCallback()(cbContext, req) + if err != nil { + t.Fatalf("beforeModel failed: %v", err) + } + + // 3. afterModel (LLM response) + resp := &model.LLMResponse{ + Content: &genai.Content{ + Role: "model", + Parts: []*genai.Part{{Text: "Hello user, calling tool now"}}, + }, + ModelVersion: "gemini-2.0-flash", + Partial: false, + } + _, err = p.AfterModelCallback()(cbContext, resp, nil) + if err != nil { + t.Fatalf("afterModel failed: %v", err) + } + + // 4. beforeTool + toolContext := &MockToolContext{ + state: mockSession.State(), + invocationID: "test-invocation", + agentName: "test_agent", + functionCallID: "call-123", + } + mockTool := &MockTool{NameVal: "test_tool"} + toolArgs := map[string]any{"param": "test"} + _, err = p.BeforeToolCallback()(toolContext, mockTool, toolArgs) + if err != nil { + t.Fatalf("beforeTool failed: %v", err) + } + + // 5. afterTool + toolResult := map[string]any{"result": "test_success"} + _, err = p.AfterToolCallback()(toolContext, mockTool, toolArgs, toolResult, nil) + if err != nil { + t.Fatalf("afterTool failed: %v", err) + } + + // 6. afterRun + p.AfterRunCallback()(invContext) + + // Verify created YAML content + filePath := filepath.Join(tempDir, "generated-recordings.yaml") + data, err := os.ReadFile(filePath) + if err != nil { + t.Fatalf("failed to read recordings file: %v", err) + } + + var recordings recording.Recordings + err = yaml.Unmarshal(data, &recordings) + if err != nil { + t.Fatalf("failed to unmarshal yaml recordings: %v", err) + } + + // Assert that raw YAML string does not contain nulls, empty collections, or default-zero/false properties + yamlStr := string(data) + if strings.Contains(yamlStr, "null") { + t.Errorf("YAML contains unexpected 'null' fields:\n%s", yamlStr) + } + if strings.Contains(yamlStr, "candidatecount") { + t.Errorf("YAML contains unexpected 'candidatecount' field:\n%s", yamlStr) + } + if strings.Contains(yamlStr, "thoughtsignature") { + t.Errorf("YAML contains unexpected 'thoughtsignature' field:\n%s", yamlStr) + } + if strings.Contains(yamlStr, "thought: false") { + t.Errorf("YAML contains unexpected 'thought: false' field:\n%s", yamlStr) + } + if strings.Contains(yamlStr, "tools:") { + t.Errorf("YAML contains unexpected 'tools:' field:\n%s", yamlStr) + } + + if len(recordings.Recordings) != 2 { + t.Fatalf("expected exactly 2 recordings, got %d", len(recordings.Recordings)) + } + + // First recording: LLM + r1 := recordings.Recordings[0] + if r1.UserMessageIndex != 0 { + t.Errorf("expected index 0, got %d", r1.UserMessageIndex) + } + if r1.AgentName != "test_agent" { + t.Errorf("expected agent 'test_agent', got %q", r1.AgentName) + } + if r1.LLMRecording == nil { + t.Fatal("expected LLM recording to be present") + } + if r1.LLMRecording.LLMRequest.Model != "gemini-2.0-flash" { + t.Errorf("expected model 'gemini-2.0-flash', got %q", r1.LLMRecording.LLMRequest.Model) + } + + // Second recording: Tool + r2 := recordings.Recordings[1] + if r2.ToolRecording == nil { + t.Fatal("expected Tool recording to be present") + } + if r2.ToolRecording.ToolCall.Name != "test_tool" { + t.Errorf("expected tool name 'test_tool', got %q", r2.ToolRecording.ToolCall.Name) + } + if r2.ToolRecording.ToolResponse.Response["result"] != "test_success" { + t.Errorf("expected response 'test_success', got %v", r2.ToolRecording.ToolResponse.Response["result"]) + } + }) + + t.Run("PathValidation", func(t *testing.T) { + tempDir := t.TempDir() + safeDir := filepath.Join(tempDir, "safe") + _ = os.Mkdir(safeDir, 0o755) + + p, mockSession, _ := setup(t, safeDir) + + tests := []struct { + name string + dir string + expectError bool + }{ + { + name: "ValidPath_InsideBaseDir", + dir: safeDir, + expectError: false, + }, + { + name: "InvalidPath_ParentTraversal", + dir: filepath.Join(safeDir, ".."), + expectError: true, + }, + { + name: "InvalidPath_AbsoluteOutside", + dir: "/etc", + expectError: true, + }, + } + + for _, tt := range tests { + t.Run(tt.name, func(t *testing.T) { + err := mockSession.State().Set("_adk_recordings_config", map[string]any{ + "dir": tt.dir, + "user_message_index": 0, + }) + if err != nil { + t.Fatalf("unexpected error: %v", err) + } + + invContext := &MockInvocationContext{ + session: mockSession, + invocationID: "test-invocation-" + tt.name, + } + + _, err = p.BeforeRunCallback()(invContext) + if tt.expectError { + if err == nil { + t.Errorf("expected path error for %q, got nil", tt.dir) + } + } else { + if err != nil { + t.Errorf("unexpected path error for %q: %v", tt.dir, err) + } + } + }) + } + }) + + t.Run("MultiTurnAppendAndDeduplication", func(t *testing.T) { + tempDir := t.TempDir() + p, mockSession, _ := setup(t, tempDir) + + // --- Turn 0 --- + err := mockSession.State().Set("_adk_recordings_config", map[string]any{ + "dir": tempDir, + "user_message_index": 0, + }) + if err != nil { + t.Fatalf("unexpected error: %v", err) + } + + invContext1 := &MockInvocationContext{session: mockSession, invocationID: "inv-0"} + _, _ = p.BeforeRunCallback()(invContext1) + + cbContext1 := &MockCallbackContext{state: mockSession.State(), invocationID: "inv-0", agentName: "test_agent"} + _, _ = p.BeforeModelCallback()(cbContext1, &model.LLMRequest{Model: "model-0"}) + _, _ = p.AfterModelCallback()(cbContext1, &model.LLMResponse{Content: &genai.Content{Parts: []*genai.Part{{Text: "Response 0"}}}, Partial: false}, nil) + p.AfterRunCallback()(invContext1) + + // --- Turn 1 Append --- + err = mockSession.State().Set("_adk_recordings_config", map[string]any{ + "dir": tempDir, + "user_message_index": 1, + }) + if err != nil { + t.Fatalf("unexpected error: %v", err) + } + + invContext2 := &MockInvocationContext{session: mockSession, invocationID: "inv-1"} + _, _ = p.BeforeRunCallback()(invContext2) + + cbContext2 := &MockCallbackContext{state: mockSession.State(), invocationID: "inv-1", agentName: "test_agent"} + _, _ = p.BeforeModelCallback()(cbContext2, &model.LLMRequest{Model: "model-1"}) + _, _ = p.AfterModelCallback()(cbContext2, &model.LLMResponse{Content: &genai.Content{Parts: []*genai.Part{{Text: "Response 1"}}}, Partial: false}, nil) + p.AfterRunCallback()(invContext2) + + // Read and verify both turns are recorded + filePath := filepath.Join(tempDir, "generated-recordings.yaml") + data, err := os.ReadFile(filePath) + if err != nil { + t.Fatalf("failed to read recordings file: %v", err) + } + + var recs recording.Recordings + _ = yaml.Unmarshal(data, &recs) + if len(recs.Recordings) != 2 { + t.Fatalf("expected 2 recordings after multi-turn run, got %d", len(recs.Recordings)) + } + if recs.Recordings[0].UserMessageIndex != 0 || recs.Recordings[1].UserMessageIndex != 1 { + t.Errorf("unexpected sequence indexes: turn 0 = %d, turn 1 = %d", recs.Recordings[0].UserMessageIndex, recs.Recordings[1].UserMessageIndex) + } + + // --- Same-Turn Deduplication/Overwrite --- + err = mockSession.State().Set("_adk_recordings_config", map[string]any{ + "dir": tempDir, + "user_message_index": 0, // Overwrite turn 0 + }) + if err != nil { + t.Fatalf("unexpected error: %v", err) + } + + invContext3 := &MockInvocationContext{session: mockSession, invocationID: "inv-0-new"} + _, _ = p.BeforeRunCallback()(invContext3) + + cbContext3 := &MockCallbackContext{state: mockSession.State(), invocationID: "inv-0-new", agentName: "test_agent"} + _, _ = p.BeforeModelCallback()(cbContext3, &model.LLMRequest{Model: "model-0"}) + _, _ = p.AfterModelCallback()(cbContext3, &model.LLMResponse{Content: &genai.Content{Parts: []*genai.Part{{Text: "Response 0 Updated"}}}, Partial: false}, nil) + p.AfterRunCallback()(invContext3) + + data, _ = os.ReadFile(filePath) + _ = yaml.Unmarshal(data, &recs) + + // Verify we still have exactly 2 recordings, but the turn 0 content is updated! + if len(recs.Recordings) != 2 { + t.Fatalf("expected exactly 2 recordings after deduplication, got %d", len(recs.Recordings)) + } + var foundTurn0 bool + for _, r := range recs.Recordings { + if r.UserMessageIndex == 0 { + foundTurn0 = true + gotText := r.LLMRecording.LLMResponses[0].Content.Parts[0].Text + if gotText != "Response 0 Updated" { + t.Errorf("expected overwritten text 'Response 0 Updated', got %q", gotText) + } + } + } + if !foundTurn0 { + t.Error("turn 0 recording was unexpectedly lost during deduplication") + } + }) + + t.Run("StreamingChunkAccumulation", func(t *testing.T) { + tempDir := t.TempDir() + p, mockSession, _ := setup(t, tempDir) + + err := mockSession.State().Set("_adk_recordings_config", map[string]any{ + "dir": tempDir, + "user_message_index": 0, + }) + if err != nil { + t.Fatalf("unexpected error: %v", err) + } + + invContext := &MockInvocationContext{session: mockSession, invocationID: "inv-stream"} + _, _ = p.BeforeRunCallback()(invContext) + + cbContext := &MockCallbackContext{state: mockSession.State(), invocationID: "inv-stream", agentName: "test_agent"} + _, _ = p.BeforeModelCallback()(cbContext, &model.LLMRequest{Model: "model-stream"}) + + // 3 Partial responses + _, _ = p.AfterModelCallback()(cbContext, &model.LLMResponse{Content: &genai.Content{Parts: []*genai.Part{{Text: "Chunk 1"}}}, Partial: true}, nil) + _, _ = p.AfterModelCallback()(cbContext, &model.LLMResponse{Content: &genai.Content{Parts: []*genai.Part{{Text: "Chunk 2"}}}, Partial: true}, nil) + _, _ = p.AfterModelCallback()(cbContext, &model.LLMResponse{Content: &genai.Content{Parts: []*genai.Part{{Text: "Chunk 3"}}}, Partial: true}, nil) + + // 1 Final non-partial response + _, _ = p.AfterModelCallback()(cbContext, &model.LLMResponse{Content: &genai.Content{Parts: []*genai.Part{{Text: "Final Chunk"}}}, Partial: false}, nil) + p.AfterRunCallback()(invContext) + + filePath := filepath.Join(tempDir, "generated-recordings.yaml") + data, _ := os.ReadFile(filePath) + var recs recording.Recordings + _ = yaml.Unmarshal(data, &recs) + + if len(recs.Recordings) != 1 { + t.Fatalf("expected 1 recording, got %d", len(recs.Recordings)) + } + llmRec := recs.Recordings[0].LLMRecording + if len(llmRec.LLMResponses) != 4 { + t.Fatalf("expected all 4 stream chunks to be accumulated, got %d", len(llmRec.LLMResponses)) + } + if llmRec.LLMResponses[3].Content.Parts[0].Text != "Final Chunk" { + t.Errorf("unexpected final response chunk text: %q", llmRec.LLMResponses[3].Content.Parts[0].Text) + } + }) + + t.Run("DisabledBypass", func(t *testing.T) { + tempDir := t.TempDir() + p, mockSession, _ := setup(t, tempDir) + + // We don't set any recording configs in mockSession state + + invContext := &MockInvocationContext{session: mockSession, invocationID: "inv-bypass"} + _, err := p.BeforeRunCallback()(invContext) + if err != nil { + t.Fatalf("beforeRun should not error on empty config, got: %v", err) + } + + cbContext := &MockCallbackContext{state: mockSession.State(), invocationID: "inv-bypass", agentName: "test_agent"} + _, err = p.BeforeModelCallback()(cbContext, &model.LLMRequest{Model: "model-bypass"}) + if err != nil { + t.Fatalf("beforeModel should not error on empty config, got: %v", err) + } + + p.AfterRunCallback()(invContext) + + filePath := filepath.Join(tempDir, "generated-recordings.yaml") + if _, err := os.Stat(filePath); err == nil { + t.Error("recordings file was unexpectedly created even though the plugin was disabled") + } + }) +} + +// --- Mock Implementations --- + +type MockState struct { + data map[string]any +} + +func (m *MockState) Set(key string, val any) error { m.data[key] = val; return nil } +func (m *MockState) Get(key string) (any, error) { return m.data[key], nil } +func (m *MockState) All() iter.Seq2[string, any] { return nil } + +type MockSession struct { + state *MockState +} + +func (m *MockSession) ID() string { return "mock-session-id" } +func (m *MockSession) AppName() string { return "mock-app" } +func (m *MockSession) UserID() string { return "mock-user" } +func (m *MockSession) State() session.State { return m.state } +func (m *MockSession) Events() session.Events { return nil } +func (m *MockSession) LastUpdateTime() time.Time { return time.Now() } + +type MockInvocationContext struct { + session *MockSession + invocationID string +} + +func (m *MockInvocationContext) Session() session.Session { return m.session } +func (m *MockInvocationContext) InvocationID() string { return m.invocationID } +func (m *MockInvocationContext) Agent() agent.Agent { return nil } +func (m *MockInvocationContext) Artifacts() agent.Artifacts { return nil } +func (m *MockInvocationContext) Memory() agent.Memory { return nil } +func (m *MockInvocationContext) Branch() string { return "" } +func (m *MockInvocationContext) UserContent() *genai.Content { return nil } +func (m *MockInvocationContext) RunConfig() *agent.RunConfig { return nil } +func (m *MockInvocationContext) LiveRequestQueue() *agent.LiveRequestQueue { return nil } +func (m *MockInvocationContext) LiveSessionResumptionHandle() string { return "" } +func (m *MockInvocationContext) SetLiveSessionResumptionHandle(handle string) {} +func (m *MockInvocationContext) EndInvocation() {} +func (m *MockInvocationContext) Ended() bool { return false } +func (m *MockInvocationContext) WithContext(ctx context.Context) agent.InvocationContext { return m } +func (m *MockInvocationContext) Value(key any) any { return nil } +func (m *MockInvocationContext) Deadline() (deadline time.Time, ok bool) { return time.Time{}, false } +func (m *MockInvocationContext) Done() <-chan struct{} { return nil } +func (m *MockInvocationContext) Err() error { return nil } + +type MockCallbackContext struct { + state session.State + invocationID string + agentName string +} + +func (m *MockCallbackContext) State() session.State { return m.state } +func (m *MockCallbackContext) ReadonlyState() session.ReadonlyState { return m.state } +func (m *MockCallbackContext) InvocationID() string { return m.invocationID } +func (m *MockCallbackContext) AgentName() string { return m.agentName } +func (m *MockCallbackContext) AppName() string { return "mock-app" } +func (m *MockCallbackContext) Branch() string { return "" } +func (m *MockCallbackContext) SessionID() string { return "mock-session-id" } +func (m *MockCallbackContext) UserID() string { return "mock-user" } +func (m *MockCallbackContext) UserContent() *genai.Content { return nil } +func (m *MockCallbackContext) Artifacts() agent.Artifacts { return nil } +func (m *MockCallbackContext) Value(key any) any { return nil } +func (m *MockCallbackContext) Deadline() (deadline time.Time, ok bool) { return time.Time{}, false } +func (m *MockCallbackContext) Done() <-chan struct{} { return nil } +func (m *MockCallbackContext) Err() error { return nil } + +type MockToolContext struct { + state session.State + invocationID string + agentName string + functionCallID string +} + +func (m *MockToolContext) State() session.State { return m.state } +func (m *MockToolContext) ReadonlyState() session.ReadonlyState { return m.state } +func (m *MockToolContext) InvocationID() string { return m.invocationID } +func (m *MockToolContext) AgentName() string { return m.agentName } +func (m *MockToolContext) FunctionCallID() string { return m.functionCallID } +func (m *MockToolContext) Actions() *session.EventActions { return nil } +func (m *MockToolContext) SearchMemory(ctx context.Context, query string) (*memory.SearchResponse, error) { + return nil, nil +} +func (m *MockToolContext) ToolConfirmation() *toolconfirmation.ToolConfirmation { return nil } +func (m *MockToolContext) RequestConfirmation(hint string, payload any) error { return nil } +func (m *MockToolContext) AppName() string { return "mock-app" } +func (m *MockToolContext) Branch() string { return "" } +func (m *MockToolContext) SessionID() string { return "mock-session-id" } +func (m *MockToolContext) UserID() string { return "mock-user" } +func (m *MockToolContext) UserContent() *genai.Content { return nil } +func (m *MockToolContext) Artifacts() agent.Artifacts { return nil } +func (m *MockToolContext) Value(key any) any { return nil } +func (m *MockToolContext) Deadline() (deadline time.Time, ok bool) { return time.Time{}, false } +func (m *MockToolContext) Done() <-chan struct{} { return nil } +func (m *MockToolContext) Err() error { return nil } + +type MockTool struct { + NameVal string +} + +func (m *MockTool) Name() string { return m.NameVal } +func (m *MockTool) Description() string { return "mock tool" } +func (m *MockTool) IsLongRunning() bool { return false } diff --git a/internal/configurable/conformance/recordplugin/yaml_clean_test.go b/internal/configurable/conformance/recordplugin/yaml_clean_test.go new file mode 100644 index 000000000..c2a6ac3e3 --- /dev/null +++ b/internal/configurable/conformance/recordplugin/yaml_clean_test.go @@ -0,0 +1,413 @@ +// Copyright 2025 Google LLC +// +// Licensed under the Apache License, Version 2.0 (the "License"); +// you may not use this file except in compliance with the License. +// You may obtain a copy of the License at +// +// http://www.apache.org/licenses/LICENSE-2.0 +// +// Unless required by applicable law or agreed to in writing, software +// distributed under the License is distributed on an "AS IS" BASIS, +// WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. +// See the License for the specific language governing permissions and +// limitations under the License. + +package recordplugin + +import ( + "testing" + + "google.golang.org/genai" + "gopkg.in/yaml.v3" +) + +func TestCleanYAMLNode(t *testing.T) { + tests := []struct { + name string + input string + expected string + }{ + { + name: "prunes null fields", + input: ` +field1: null +field2: ~ +field3: "value" +`, + expected: `field3: "value" +`, + }, + { + name: "prunes default zero values and false booleans outside tool args", + input: ` +candidate_count: 0 +max_output_tokens: 0 +thought: false +text: "keep me" +`, + expected: `text: "keep me" +`, + }, + { + name: "retains zero values and false booleans inside tool args", + input: ` +tool_call: + name: "my_tool" + args: + param1: 0 + param2: false +`, + expected: `tool_call: + name: "my_tool" + args: + param1: 0 + param2: false +`, + }, + { + name: "retains zero values and false booleans inside tool responses", + input: ` +tool_response: + response: + success: false + count: 0 +`, + expected: `tool_response: + response: + success: false + count: 0 +`, + }, + { + name: "prunes empty sequences but retains empty mappings", + input: ` +thought_signature: [] +google_search: {} +text: "test" +`, + expected: `google_search: {} +text: "test" +`, + }, + { + name: "retains thought_signature base64 scalar as string", + input: ` +text: "test" +thought_signature: "SGVsbG8=" +`, + expected: `text: "test" +thought_signature: "SGVsbG8=" +`, + }, + { + name: "retains thoughtsignature base64 scalar as string", + input: ` +text: "test" +thoughtsignature: "SGVsbG8=" +`, + expected: `text: "test" +thought_signature: "SGVsbG8=" +`, + }, + { + name: "converts thought_signature sequence of integers to base64 string", + input: ` +text: "test" +thought_signature: + - 72 + - 101 + - 108 + - 108 + - 111 +`, + expected: `text: "test" +thought_signature: "SGVsbG8=" +`, + }, + { + name: "converts thoughtsignature sequence of integers to base64 string", + input: ` +text: "test" +thoughtsignature: + - 72 + - 101 + - 108 + - 108 + - 111 +`, + expected: `text: "test" +thought_signature: "SGVsbG8=" +`, + }, + { + name: "deep nested pruning", + input: ` +recordings: + - user_message_index: 0 + agent_name: "search" + llm_recording: + llm_request: + model: "gemini" + contents: + - parts: + - media_resolution: null + text: "hello" + thought_signature: [] +`, + expected: `recordings: + - user_message_index: 0 + agent_name: "search" + llm_recording: + llm_request: + model: "gemini" + contents: + - parts: + - text: "hello" +`, + }, + { + name: "cleans bloated parameters schema and translates keys", + input: ` +functiondeclarations: + - name: "validate_email" + description: "checks email" + parametersjsonschema: + id: "" + schema: "" + ref: "" + comment: "" + defs: {} + definitions: {} + dependentrequired: {} + extra: {} + type: "object" + required: + - "email" + properties: + email: + id: "" + type: "string" + description: "the email address" + dependentrequired: {} + responsejsonschema: + type: "boolean" +`, + expected: ` +function_declarations: + - name: "validate_email" + description: "checks email" + parameters_json_schema: + title: "validate_emailParams" + type: "object" + required: + - "email" + properties: + email: + title: "Email" + type: "string" + description: "the email address" +`, + }, + { + name: "cleans schema parameters mapping additionalproperties and propertyorder, and pruning unsupported openapi properties", + input: ` +functiondeclarations: + - name: "validate_email" + parametersjsonschema: + type: "object" + required: + - "email" + properties: + email: + type: "string" + additionalproperties: + not: {} + propertyorder: + - "email" + deprecated: false + readonly: false + writeonly: false + uniqueitems: false +`, + expected: ` +function_declarations: + - name: "validate_email" + parameters_json_schema: + title: "validate_emailParams" + type: "object" + required: + - "email" + properties: + email: + title: "Email" + type: "string" +`, + }, + { + name: "prunes empty strings for general fields but retains them inside tool args", + input: ` +llm_recording: + llm_responses: + - content: + parts: + - text: "done" + session_resumption_handle: "" + error_code: "" + error_message: "" +tool_call: + name: "my_tool" + args: + param_empty_str: "" +`, + expected: ` +llm_recording: + llm_responses: + - content: + parts: + - text: "done" +tool_call: + name: "my_tool" + args: + param_empty_str: "" +`, + }, + { + name: "prunes empty GenerateContentConfig and google search properties", + input: ` +llm_recording: + llm_request: + config: + response_mime_type: "" + cached_content: "" + media_resolution: "" + service_tier: "" + tools: + - google_search: + blocking_confidence: "" + labels: {} +`, + expected: ` +llm_recording: + llm_request: + config: + tools: + - google_search: {} +`, + }, + { + name: "normalizes raw string system_instruction to structured object", + input: ` +llm_request: + config: + system_instruction: "You are a booking assistant." +`, + expected: ` +llm_request: + config: + system_instruction: "You are a booking assistant." +`, + }, + { + name: "injects part_metadata to contents parts", + input: ` +llm_request: + contents: + - role: "user" + parts: + - text: "Hello" +`, + expected: ` +llm_request: + contents: + - role: "user" + parts: + - text: "Hello" +`, + }, + { + name: "preserves existing part_metadata if already present", + input: ` +llm_request: + contents: + - role: "user" + parts: + - text: "Hello" + part_metadata: + some_key: "some_value" +`, + expected: ` +llm_request: + contents: + - role: "user" + parts: + - text: "Hello" + part_metadata: + some_key: "some_value" +`, + }, + } + + for _, tc := range tests { + t.Run(tc.name, func(t *testing.T) { + var node yaml.Node + err := yaml.Unmarshal([]byte(tc.input), &node) + if err != nil { + t.Fatalf("failed to unmarshal test input: %v", err) + } + + cleanYAMLNode(&node, false) + + output, err := yaml.Marshal(&node) + if err != nil { + t.Fatalf("failed to marshal cleaned node: %v", err) + } + + gotClean := normalizeYAML(string(output)) + expectedClean := normalizeYAML(tc.expected) + if gotClean != expectedClean { + t.Errorf("mismatch.\nGot:\n%s\nExpected:\n%s", gotClean, expectedClean) + } + }) + } +} + +func normalizeYAML(s string) string { + var node yaml.Node + if err := yaml.Unmarshal([]byte(s), &node); err != nil { + return s + } + out, _ := yaml.Marshal(&node) + return string(out) +} + +func TestFunctionDeclarationTitleInjection(t *testing.T) { + decl := &genai.FunctionDeclaration{ + Name: "validate_email", + Description: "checks email", + Parameters: &genai.Schema{ + Type: genai.TypeObject, + Required: []string{"email"}, + Properties: map[string]*genai.Schema{ + "email": { + Type: genai.TypeString, + Description: "the email address", + }, + }, + }, + } + + var node yaml.Node + err := node.Encode(decl) + if err != nil { + t.Fatalf("failed to encode: %v", err) + } + + cleanYAMLNode(&node, false) + + output, err := yaml.Marshal(&node) + if err != nil { + t.Fatalf("failed to marshal: %v", err) + } + + t.Logf("Serialized:\n%s", string(output)) +} diff --git a/internal/configurable/conformance/replayplugin/recording/recording.go b/internal/configurable/conformance/replayplugin/recording/recording.go index 98eaed777..99474b109 100644 --- a/internal/configurable/conformance/replayplugin/recording/recording.go +++ b/internal/configurable/conformance/replayplugin/recording/recording.go @@ -29,18 +29,18 @@ type Recordings struct { // Recording represents a single interaction recording, ordered by request timestamp. type Recording struct { // Index of the user message this recording belongs to (0-based). - UserMessageIndex int `yaml:"usermessageindex"` + UserMessageIndex int `yaml:"user_message_index"` // Name of the agent. - AgentName string `yaml:"agentname"` + AgentName string `yaml:"agent_name"` // oneof fields - start // LLM request-response pair. - LLMRecording *LLMRecording `yaml:"llmrecording,omitempty"` + LLMRecording *LLMRecording `yaml:"llm_recording,omitempty"` // Tool call-response pair. - ToolRecording *ToolRecording `yaml:"toolrecording,omitempty"` + ToolRecording *ToolRecording `yaml:"tool_recording,omitempty"` // oneof fields - end @@ -51,17 +51,17 @@ type Recording struct { // LLMRecording represents a paired LLM request and response. type LLMRecording struct { // Required. The LLM request. - LLMRequest *model.LLMRequest `yaml:"llmrequest,omitempty"` + LLMRequest *model.LLMRequest `yaml:"llm_request,omitempty"` // Required. The LLM response. - LLMResponses []*model.LLMResponse `yaml:"llmresponses,omitempty"` + LLMResponses []*model.LLMResponse `yaml:"llm_responses,omitempty"` } // ToolRecording represents a paired tool call and response. type ToolRecording struct { // Required. The tool call. - ToolCall *genai.FunctionCall `yaml:"toolcall,omitempty"` + ToolCall *genai.FunctionCall `yaml:"tool_call,omitempty"` // Required. The tool response. - ToolResponse *genai.FunctionResponse `yaml:"toolresponse,omitempty"` + ToolResponse *genai.FunctionResponse `yaml:"tool_response,omitempty"` } diff --git a/internal/configurable/conformance/replayplugin/replay_plugin.go b/internal/configurable/conformance/replayplugin/replay_plugin.go index ffda62baf..4efeb1f52 100644 --- a/internal/configurable/conformance/replayplugin/replay_plugin.go +++ b/internal/configurable/conformance/replayplugin/replay_plugin.go @@ -129,11 +129,15 @@ func (p *replayPlugin) beforeModel(ctx agent.CallbackContext, req *model.LLMRequ return nil, err } + if len(recording.LLMResponses) == 0 { + return nil, fmt.Errorf("no LLM responses found in recording for agent %q", agentName) + } + return recording.LLMResponses[0], nil } // beforeTool intercepts tool calls, verifies them against the recording, and returns the recorded response. -func (p *replayPlugin) beforeTool(ctx tool.Context, t tool.Tool, args map[string]any) (map[string]any, error) { +func (p *replayPlugin) beforeTool(ctx agent.ToolContext, t tool.Tool, args map[string]any) (map[string]any, error) { on, err := p.isReplayModeOn(ctx.State()) if err != nil { return nil, err @@ -312,8 +316,7 @@ func (p *replayPlugin) loadInvocationState(ctx agent.InvocationContext) (*invoca return nil, fmt.Errorf("failed to parse recordings from %s: %w", recordingsPath, err) } - removeUnderscores(&root) - fixTypeMismatches(&root) + normalizeYAMLNode(&root) var recordings recording.Recordings if err := root.Decode(&recordings); err != nil { @@ -425,6 +428,17 @@ func verifyLLMRequestMatch(expectedLLMRequest, actualLLMRequest *model.LLMReques cmpopts.EquateEmpty(), } + for _, toolAny := range expectedLLMRequest.Tools { + if funcDecl, ok := toolAny.(*genai.FunctionDeclaration); ok { + funcDecl.Description = normalizeDescription(funcDecl.Description) + } + } + for _, toolAny := range actualLLMRequest.Tools { + if funcDecl, ok := toolAny.(*genai.FunctionDeclaration); ok { + funcDecl.Description = normalizeDescription(funcDecl.Description) + } + } + if transferToolAny, ok := expectedLLMRequest.Tools["transfer_to_agent"]; ok { transferTool := transferToolAny.(*genai.FunctionDeclaration) transferTool.Description = `Transfer the question to another agent. @@ -434,6 +448,7 @@ This tool hands off control to another agent when it's more suitable to answer t if expectedLLMRequest.Config != nil { for _, tool := range expectedLLMRequest.Config.Tools { for _, funcDecl := range tool.FunctionDeclarations { + funcDecl.Description = normalizeDescription(funcDecl.Description) if funcDecl.Name == "transfer_to_agent" { funcDecl.Description = `Transfer the question to another agent. This tool hands off control to another agent when it's more suitable to answer the user's question according to the agent's description.` @@ -442,6 +457,14 @@ This tool hands off control to another agent when it's more suitable to answer t } } + if actualLLMRequest.Config != nil { + for _, tool := range actualLLMRequest.Config.Tools { + for _, funcDecl := range tool.FunctionDeclarations { + funcDecl.Description = normalizeDescription(funcDecl.Description) + } + } + } + // Compare! // cmp.Diff returns an empty string if they are equal, otherwise a human-readable diff. if diff := cmp.Diff(expectedLLMRequest, actualLLMRequest, opts...); diff != "" { @@ -580,3 +603,19 @@ func verifyToolCallMatch(expectedToolCall *genai.FunctionCall, toolName string, return nil } + +func normalizeDescription(desc string) string { + lines := strings.Split(desc, "\n") + var cleaned []string + for _, line := range lines { + cleaned = append(cleaned, strings.TrimSpace(line)) + } + // Remove empty lines at the start and end + for len(cleaned) > 0 && cleaned[0] == "" { + cleaned = cleaned[1:] + } + for len(cleaned) > 0 && cleaned[len(cleaned)-1] == "" { + cleaned = cleaned[:len(cleaned)-1] + } + return strings.Join(cleaned, "\n") +} diff --git a/internal/configurable/conformance/replayplugin/replay_plugin_internal_test.go b/internal/configurable/conformance/replayplugin/replay_plugin_internal_test.go new file mode 100644 index 000000000..8701d868b --- /dev/null +++ b/internal/configurable/conformance/replayplugin/replay_plugin_internal_test.go @@ -0,0 +1,80 @@ +// Copyright 2026 Google LLC +// +// Licensed under the Apache License, Version 2.0 (the "License"); +// you may not use this file except in compliance with the License. +// You may obtain a copy of the License at +// +// http://www.apache.org/licenses/LICENSE-2.0 +// +// Unless required by applicable law or agreed to in writing, software +// distributed under the License is distributed on an "AS IS" BASIS, +// WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. +// See the License for the specific language governing permissions and +// limitations under the License. + +package replayplugin + +import "testing" + +func TestNormalizeDescription(t *testing.T) { + tests := []struct { + name string + input string + expected string + }{ + { + name: "Simple single line", + input: "Simple description", + expected: "Simple description", + }, + { + name: "Single line with leading and trailing spaces", + input: " Simple description ", + expected: "Simple description", + }, + { + name: "Multiple lines with indentation", + input: " Line 1 \n Line 2 ", + expected: "Line 1\nLine 2", + }, + { + name: "Empty lines at start and end", + input: "\n\nLine 1\nLine 2\n\n", + expected: "Line 1\nLine 2", + }, + { + name: "Whitespace-only lines at start and end", + input: " \n \nLine 1\nLine 2\n \n ", + expected: "Line 1\nLine 2", + }, + { + name: "Empty lines in the middle are preserved", + input: "Line 1\n\nLine 2", + expected: "Line 1\n\nLine 2", + }, + { + name: "Empty input", + input: "", + expected: "", + }, + { + name: "Whitespace-only input", + input: " \n \n ", + expected: "", + }, + { + name: "Complex formatting with tabs and newlines", + input: "\t\n Line 1 \t \n\tLine 2\n \t\n", + expected: "Line 1\nLine 2", + }, + } + + for _, tt := range tests { + t.Run(tt.name, func(t *testing.T) { + actual := normalizeDescription(tt.input) + if actual != tt.expected { + t.Errorf("normalizeDescription(%q) = %q; want %q", tt.input, actual, tt.expected) + } + }) + } +} diff --git a/internal/configurable/conformance/replayplugin/replay_plugin_test.go b/internal/configurable/conformance/replayplugin/replay_plugin_test.go index bc2bfb37e..785b91c63 100644 --- a/internal/configurable/conformance/replayplugin/replay_plugin_test.go +++ b/internal/configurable/conformance/replayplugin/replay_plugin_test.go @@ -121,6 +121,83 @@ recordings: } }) + t.Run("BeforeModelCallback_WithSingleLlmResponse_ReturnsRecordedResponse", func(t *testing.T) { + plugin, mockSession, _ := setup(t) + tempDir := t.TempDir() + + // 1. Create recording file with singular llm_response instead of plural llm_responses + recordingsYaml := ` +recordings: + - user_message_index: 0 + agent_name: "test_agent" + llm_recording: + llm_request: + model: "gemini-2.0-flash" + contents: + - role: "user" + parts: + - text: "Hello" + llm_response: + content: + role: "model" + parts: + - text: "Recorded response" +` + createRecordingsFile(t, tempDir, recordingsYaml) + + // 2. Setup replay config + err := mockSession.State().Set("_adk_replay_config", map[string]any{ + "dir": tempDir, + "user_message_index": 0, + }) + if err != nil { + t.Fatalf("unexpected error: %v", err) + } + + // 3. Load recordings (BeforeRunCallback) + invContext := &MockInvocationContext{ + session: mockSession, + invocationID: "test-invocation", + } + _, err = plugin.BeforeRunCallback()(invContext) + if err != nil { + t.Fatalf("unexpected error: %v", err) + } + + // 4. Call BeforeModelCallback with matching request + cbContext := &MockCallbackContext{ + state: mockSession.State(), + invocationID: "test-invocation", + agentName: "test_agent", + } + + request := &model.LLMRequest{ + Model: "gemini-2.0-flash", + Contents: []*genai.Content{ + { + Role: "user", + Parts: []*genai.Part{{Text: "Hello"}}, + }, + }, + } + + result, err := plugin.BeforeModelCallback()(cbContext, request) + // 5. Verify + if err != nil { + t.Fatalf("unexpected error: %v", err) + } + if result == nil { + t.Fatal("expected non-nil result") + } + if result.Content == nil { + t.Fatal("expected non-nil result.Content") + } + + if got := result.Content.Parts[0].Text; got != "Recorded response" { + t.Errorf("expected %q, got %q", "Recorded response", got) + } + }) + t.Run("BeforeModelCallback_RequestMismatch_ReturnsEmpty", func(t *testing.T) { plugin, mockSession, _ := setup(t) tempDir := t.TempDir() @@ -397,12 +474,12 @@ func (m *MockInvocationContext) Memory() agent.Memory func (m *MockInvocationContext) Branch() string { return "" } func (m *MockInvocationContext) UserContent() *genai.Content { return nil } func (m *MockInvocationContext) RunConfig() *agent.RunConfig { return nil } // Use context? No, RunConfig struct. +func (m *MockInvocationContext) LiveRequestQueue() *agent.LiveRequestQueue { return nil } +func (m *MockInvocationContext) LiveSessionResumptionHandle() string { return "" } +func (m *MockInvocationContext) SetLiveSessionResumptionHandle(handle string) {} func (m *MockInvocationContext) EndInvocation() {} func (m *MockInvocationContext) Ended() bool { return false } func (m *MockInvocationContext) WithContext(ctx context.Context) agent.InvocationContext { return m } -func (m *MockInvocationContext) LiveRequestQueue() *agent.LiveRequestQueue { return nil } -func (m *MockInvocationContext) SetLiveSessionResumptionHandle(_ string) {} -func (m *MockInvocationContext) LiveSessionResumptionHandle() string { return "" } func (m *MockInvocationContext) Value(key any) any { return nil } func (m *MockInvocationContext) Deadline() (deadline time.Time, ok bool) { return time.Time{}, false } func (m *MockInvocationContext) Done() <-chan struct{} { return nil } diff --git a/internal/configurable/conformance/replayplugin/yaml_utils.go b/internal/configurable/conformance/replayplugin/yaml_utils.go index ac3ec334c..f344f23ce 100644 --- a/internal/configurable/conformance/replayplugin/yaml_utils.go +++ b/internal/configurable/conformance/replayplugin/yaml_utils.go @@ -20,47 +20,41 @@ import ( "gopkg.in/yaml.v3" ) -var toIgnore = map[string]struct{}{"thought_signature": {}, "http_options": {}, "args": {}, "response": {}} +var toIgnore = map[string]struct{}{ + "thought_signature": {}, + "http_options": {}, + "args": {}, + "response": {}, +} + +var toIgnoreButRecurse = map[string]struct{}{ + "user_message_index": {}, + "agent_name": {}, + "llm_recording": {}, + "llm_request": {}, + "llm_responses": {}, + "tool_recording": {}, + "tool_call": {}, + "tool_response": {}, +} -func removeUnderscores(node *yaml.Node) { +// normalizeYAMLNode standardizes the recorded YAML AST in a single tree traversal. +// It strips underscores from structural Go field keys to match case-insensitive fields, +// while preserving user freeform tool arguments/responses, and resolves type mismatches. +func normalizeYAMLNode(node *yaml.Node) { switch node.Kind { case yaml.DocumentNode, yaml.SequenceNode: - // If it's a document or a list, just pass through to its children for _, child := range node.Content { - removeUnderscores(child) + normalizeYAMLNode(child) } case yaml.MappingNode: - // A MappingNode's content is a flat array alternating [key, value, key, value...] for i := 0; i < len(node.Content); i += 2 { keyNode := node.Content[i] valueNode := node.Content[i+1] - if _, ok := toIgnore[keyNode.Value]; ok { - continue - } - - // Strip the underscore from the key (e.g., "first_name" -> "firstname") - keyNode.Value = strings.ReplaceAll(keyNode.Value, "_", "") - - // Continue walking down into the value in case of nested objects - removeUnderscores(valueNode) - } - } -} - -func fixTypeMismatches(n *yaml.Node) { - switch n.Kind { - case yaml.DocumentNode, yaml.SequenceNode: - for _, child := range n.Content { - fixTypeMismatches(child) - } - - case yaml.MappingNode: - for i := 0; i < len(n.Content); i += 2 { - keyNode := n.Content[i] - valueNode := n.Content[i+1] + // 1. Fix known type mismatches and singular/plural fields switch keyNode.Value { - case "systeminstruction": + case "systeminstruction", "system_instruction": if valueNode.Kind == yaml.ScalarNode { val := valueNode.Value valueNode.Kind = yaml.MappingNode @@ -86,10 +80,33 @@ func fixTypeMismatches(n *yaml.Node) { }, } } + case "llmresponse", "llm_response": + if valueNode.Kind == yaml.MappingNode { + origCopy := *valueNode + valueNode.Kind = yaml.SequenceNode + valueNode.Tag = "!!seq" + valueNode.Value = "" + valueNode.Content = []*yaml.Node{&origCopy} + } + keyNode.Value = "llm_responses" + } + + // 2. Skip processing/recursion for ignored freeform keys (like tool args or response payload) + if _, ok := toIgnore[keyNode.Value]; ok { + continue } - // Recurse into the value to catch nested structures - fixTypeMismatches(valueNode) + // 3. For standard metadata fields, preserve snake_case but recurse into value + if _, ok := toIgnoreButRecurse[keyNode.Value]; ok { + normalizeYAMLNode(valueNode) + continue + } + + // 4. Strip underscores to match camelCase Go struct field names case-insensitively + keyNode.Value = strings.ReplaceAll(keyNode.Value, "_", "") + + // 5. Recurse to process nested fields + normalizeYAMLNode(valueNode) } } } diff --git a/internal/configurable/conformance/replayplugin/yaml_utils_test.go b/internal/configurable/conformance/replayplugin/yaml_utils_test.go new file mode 100644 index 000000000..e36336596 --- /dev/null +++ b/internal/configurable/conformance/replayplugin/yaml_utils_test.go @@ -0,0 +1,173 @@ +// Copyright 2026 Google LLC +// +// Licensed under the Apache License, Version 2.0 (the "License"); +// you may not use this file except in compliance with the License. +// You may obtain a copy of the License at +// +// http://www.apache.org/licenses/LICENSE-2.0 +// +// Unless required by applicable law or agreed to in writing, software +// distributed under the License is distributed on an "AS IS" BASIS, +// WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. +// See the License for the specific language governing permissions and +// limitations under the License. + +package replayplugin + +import ( + "testing" + + "gopkg.in/yaml.v3" +) + +func TestRemoveUnderscores(t *testing.T) { + tests := []struct { + name string + input string + expected string + }{ + { + name: "strips underscores from standard keys", + input: ` +first_name: "John" +last_name: "Doe" +`, + expected: ` +firstname: "John" +lastname: "Doe" +`, + }, + { + name: "ignores fields specified in toIgnore", + input: ` +thought_signature: "SGVsbG8=" +http_options: "options" +args: "arguments" +response: "result" +`, + expected: ` +thought_signature: "SGVsbG8=" +http_options: "options" +args: "arguments" +response: "result" +`, + }, + { + name: "ignores but recurses fields specified in toIgnoreButRecurse", + input: ` +user_message_index: 0 +agent_name: "test" +llm_recording: + nested_field_with_underscore: "nested" +`, + expected: ` +user_message_index: 0 +agent_name: "test" +llm_recording: + nestedfieldwithunderscore: "nested" +`, + }, + } + + for _, tc := range tests { + t.Run(tc.name, func(t *testing.T) { + var node yaml.Node + err := yaml.Unmarshal([]byte(tc.input), &node) + if err != nil { + t.Fatalf("Unmarshal failed: %v", err) + } + + normalizeYAMLNode(&node) + + output, err := yaml.Marshal(&node) + if err != nil { + t.Fatalf("Marshal failed: %v", err) + } + + gotClean := normalizeYAML(string(output)) + expectedClean := normalizeYAML(tc.expected) + if gotClean != expectedClean { + t.Errorf("mismatch.\nGot:\n%s\nExpected:\n%s", gotClean, expectedClean) + } + }) + } +} + +func TestFixTypeMismatches(t *testing.T) { + tests := []struct { + name string + input string + expected string + }{ + { + name: "wraps llmresponse mapping into sequence llm_responses", + input: ` +llmresponse: + content: + role: "model" +`, + expected: ` +llm_responses: + - content: + role: "model" +`, + }, + { + name: "wraps llm_response mapping into sequence llm_responses", + input: ` +llm_response: + content: + role: "model" +`, + expected: ` +llm_responses: + - content: + role: "model" +`, + }, + { + name: "normalizes systeminstruction scalar string to structured object", + input: ` +systeminstruction: "You are a helpful assistant." +`, + expected: ` +systeminstruction: + role: user + parts: + - text: You are a helpful assistant. +`, + }, + } + + for _, tc := range tests { + t.Run(tc.name, func(t *testing.T) { + var node yaml.Node + err := yaml.Unmarshal([]byte(tc.input), &node) + if err != nil { + t.Fatalf("Unmarshal failed: %v", err) + } + + normalizeYAMLNode(&node) + + output, err := yaml.Marshal(&node) + if err != nil { + t.Fatalf("Marshal failed: %v", err) + } + + gotClean := normalizeYAML(string(output)) + expectedClean := normalizeYAML(tc.expected) + if gotClean != expectedClean { + t.Errorf("mismatch.\nGot:\n%s\nExpected:\n%s", gotClean, expectedClean) + } + }) + } +} + +func normalizeYAML(s string) string { + var node yaml.Node + if err := yaml.Unmarshal([]byte(s), &node); err != nil { + return s + } + out, _ := yaml.Marshal(&node) + return string(out) +} diff --git a/internal/context/callback_context.go b/internal/context/callback_context.go index ca56ff0ae..251a507f3 100644 --- a/internal/context/callback_context.go +++ b/internal/context/callback_context.go @@ -15,111 +15,22 @@ package context import ( - "context" - "iter" - - "google.golang.org/genai" - "google.golang.org/adk/agent" - "google.golang.org/adk/artifact" "google.golang.org/adk/session" ) -type internalArtifacts struct { - agent.Artifacts - eventActions *session.EventActions -} - -func (ia *internalArtifacts) Save(ctx context.Context, name string, data *genai.Part) (*artifact.SaveResponse, error) { - resp, err := ia.Artifacts.Save(ctx, name, data) - if err != nil { - return resp, err - } - if ia.eventActions != nil { - if ia.eventActions.ArtifactDelta == nil { - ia.eventActions.ArtifactDelta = make(map[string]int64) - } - // TODO: RWLock, check the version stored is newer in case multiple tools save the same file. - ia.eventActions.ArtifactDelta[name] = resp.Version - } - return resp, nil -} - +// NewCallbackContext returns a CallbackContext suitable for model, tool, and +// related callbacks. The returned context's Artifacts().Save tracks each saved +// artifact's version into the underlying EventActions.ArtifactDelta. func NewCallbackContext(ctx agent.InvocationContext) agent.CallbackContext { - return newCallbackContext(ctx, make(map[string]any), make(map[string]int64)) + return agent.NewCallbackContextWithArtifactTracking(ctx, nil) } +// NewCallbackContextWithDelta returns a CallbackContext that uses the given +// stateDelta and artifactDelta maps as the initial backing storage for its +// EventActions. The returned context's Artifacts().Save tracks each saved +// artifact's version into artifactDelta. func NewCallbackContextWithDelta(ctx agent.InvocationContext, stateDelta map[string]any, artifactDelta map[string]int64) agent.CallbackContext { - return newCallbackContext(ctx, stateDelta, artifactDelta) -} - -func newCallbackContext(ctx agent.InvocationContext, stateDelta map[string]any, artifactDelta map[string]int64) *callbackContext { - rCtx := NewReadonlyContext(ctx) - eventActions := &session.EventActions{StateDelta: stateDelta, ArtifactDelta: artifactDelta} - return &callbackContext{ - ReadonlyContext: rCtx, - invocationCtx: ctx, - eventActions: eventActions, - artifacts: &internalArtifacts{ - Artifacts: ctx.Artifacts(), - eventActions: eventActions, - }, - } -} - -// TODO: unify with agent.callbackContext - -type callbackContext struct { - agent.ReadonlyContext - artifacts *internalArtifacts - invocationCtx agent.InvocationContext - eventActions *session.EventActions -} - -func (c *callbackContext) Artifacts() agent.Artifacts { - return c.artifacts -} - -func (c *callbackContext) AgentName() string { - return c.invocationCtx.Agent().Name() -} - -func (c *callbackContext) ReadonlyState() session.ReadonlyState { - return c.invocationCtx.Session().State() -} - -func (c *callbackContext) State() session.State { - return &callbackContextState{ctx: c} -} - -func (c *callbackContext) InvocationID() string { - return c.invocationCtx.InvocationID() -} - -func (c *callbackContext) UserContent() *genai.Content { - return c.invocationCtx.UserContent() -} - -type callbackContextState struct { - ctx *callbackContext -} - -func (c *callbackContextState) Get(key string) (any, error) { - if c.ctx.eventActions != nil && c.ctx.eventActions.StateDelta != nil { - if val, ok := c.ctx.eventActions.StateDelta[key]; ok { - return val, nil - } - } - return c.ctx.invocationCtx.Session().State().Get(key) -} - -func (c *callbackContextState) Set(key string, val any) error { - if c.ctx.eventActions != nil && c.ctx.eventActions.StateDelta != nil { - c.ctx.eventActions.StateDelta[key] = val - } - return c.ctx.invocationCtx.Session().State().Set(key, val) -} - -func (c *callbackContextState) All() iter.Seq2[string, any] { - return c.ctx.invocationCtx.Session().State().All() + actions := &session.EventActions{StateDelta: stateDelta, ArtifactDelta: artifactDelta} + return agent.NewCallbackContextWithArtifactTracking(ctx, actions) } diff --git a/internal/context/context_test.go b/internal/context/context_test.go index 8a18e20ea..6e615a294 100644 --- a/internal/context/context_test.go +++ b/internal/context/context_test.go @@ -16,12 +16,15 @@ package context import ( "context" + "strings" "testing" "github.com/google/go-cmp/cmp" "github.com/google/go-cmp/cmp/cmpopts" + "github.com/google/uuid" "google.golang.org/adk/agent" + "google.golang.org/adk/platform" ) func TestReadonlyContext(t *testing.T) { @@ -64,3 +67,53 @@ func TestWithContext(t *testing.T) { t.Errorf("WithContext() mismatch (-want +got):\n%s", diff) } } + +func TestNewInvocationContextGeneratesIDWithProvider(t *testing.T) { + ctx := platform.WithUUIDProvider(t.Context(), func() string { return "fixed" }) + inv := NewInvocationContext(ctx, InvocationContextParams{}) + + if got, want := inv.InvocationID(), "e-fixed"; got != want { + t.Errorf("InvocationID() = %q, want %q", got, want) + } +} + +func TestNewInvocationContextRespectsExplicitID(t *testing.T) { + // An explicit InvocationID must be used verbatim, leaving the provider unused. + ctx := platform.WithUUIDProvider(t.Context(), func() string { return "fixed" }) + inv := NewInvocationContext(ctx, InvocationContextParams{InvocationID: "explicit"}) + + if got, want := inv.InvocationID(), "explicit"; got != want { + t.Errorf("InvocationID() = %q, want %q", got, want) + } +} + +func TestNewInvocationContextDefaultID(t *testing.T) { + inv := NewInvocationContext(t.Context(), InvocationContextParams{}) + + id := inv.InvocationID() + rest, ok := strings.CutPrefix(id, "e-") + if !ok { + t.Fatalf("InvocationID() = %q, want \"e-\" prefix", id) + } + if _, err := uuid.Parse(rest); err != nil { + t.Errorf("InvocationID() = %q, suffix not a valid UUID: %v", id, err) + } +} + +func TestInvocationContext_LiveSessionResumptionHandle(t *testing.T) { + inv := NewInvocationContext(t.Context(), InvocationContextParams{}) + + iCtx, ok := inv.(*InvocationContext) + if !ok { + t.Fatalf("NewInvocationContext did not return *InvocationContext") + } + + if iCtx.LiveSessionResumptionHandle() != "" { + t.Errorf("expected empty handle, got %q", iCtx.LiveSessionResumptionHandle()) + } + + iCtx.SetLiveSessionResumptionHandle("test-handle") + if iCtx.LiveSessionResumptionHandle() != "test-handle" { + t.Errorf("expected test-handle, got %q", iCtx.LiveSessionResumptionHandle()) + } +} diff --git a/internal/context/invocation_context.go b/internal/context/invocation_context.go index aa66e94d3..b1bd1acbf 100644 --- a/internal/context/invocation_context.go +++ b/internal/context/invocation_context.go @@ -17,10 +17,10 @@ package context import ( "context" - "github.com/google/uuid" "google.golang.org/genai" "google.golang.org/adk/agent" + "google.golang.org/adk/platform" "google.golang.org/adk/session" ) @@ -42,7 +42,7 @@ type InvocationContextParams struct { func NewInvocationContext(ctx context.Context, params InvocationContextParams) agent.InvocationContext { if params.InvocationID == "" { - params.InvocationID = "e-" + uuid.NewString() + params.InvocationID = "e-" + platform.NewUUID(ctx) } return &InvocationContext{ Context: ctx, @@ -92,14 +92,6 @@ func (c *InvocationContext) LiveRequestQueue() *agent.LiveRequestQueue { return c.params.LiveRequestQueue } -func (c *InvocationContext) SetLiveSessionResumptionHandle(handle string) { - c.params.LiveSessionResumptionHandle = handle -} - -func (c *InvocationContext) LiveSessionResumptionHandle() string { - return c.params.LiveSessionResumptionHandle -} - func (c *InvocationContext) EndInvocation() { c.params.EndInvocation = true } @@ -108,6 +100,18 @@ func (c *InvocationContext) Ended() bool { return c.params.EndInvocation } +// LiveSessionResumptionHandle returns the active live session's resumption handle. +// This handle is used to reconnect and resume a bidirectional streaming session with the model. +func (c *InvocationContext) LiveSessionResumptionHandle() string { + return c.params.LiveSessionResumptionHandle +} + +// SetLiveSessionResumptionHandle stores the latest resumption handle received from the model. +// This allows subsequent reconnection attempts in the live loop to resume the same session state. +func (c *InvocationContext) SetLiveSessionResumptionHandle(handle string) { + c.params.LiveSessionResumptionHandle = handle +} + func (c *InvocationContext) WithContext(ctx context.Context) agent.InvocationContext { newCtx := *c newCtx.Context = ctx diff --git a/internal/llminternal/agent_transfer.go b/internal/llminternal/agent_transfer.go index ee021b88c..a58449d3b 100644 --- a/internal/llminternal/agent_transfer.go +++ b/internal/llminternal/agent_transfer.go @@ -65,6 +65,7 @@ import ( // // TODO: implement it in the runners package and update this doc. +// AgentTransferRequestProcessor processes agent transfer requests. func AgentTransferRequestProcessor(ctx agent.InvocationContext, req *model.LLMRequest, f *Flow) iter.Seq2[*session.Event, error] { return func(yield func(*session.Event, error) bool) { // TODO: support agent types other than LLMAgent, that have parent/subagents? @@ -82,13 +83,12 @@ func AgentTransferRequestProcessor(ctx agent.InvocationContext, req *model.LLMRe // TODO(hyangah): why do we set this up in request processor // instead of registering this as a normal function tool of the Agent? - transferToAgentTool := &TransferToAgentTool{} - si, err := instructionsForTransferToAgent(agent, parents[agent.Name()], targets, transferToAgentTool) + transferToAgentTool, err := NewTransferToAgentTool(agent, parents[agent.Name()], targets) if err != nil { yield(nil, err) return } - utils.AppendInstructions(req, si) + utils.AppendInstructions(req, transferToAgentTool.instructions) err = appendTools(req, transferToAgentTool) if err != nil { yield(nil, err) @@ -96,7 +96,25 @@ func AgentTransferRequestProcessor(ctx agent.InvocationContext, req *model.LLMRe } } -type TransferToAgentTool struct{} +const transferAgentName = "transfer_to_agent" + +// TransferToAgentTool is a tool that handles transferring control to another agent. +type TransferToAgentTool struct { + instructions string + supportedAgents []agent.Agent +} + +// NewTransferToAgentTool creates a new TransferToAgentTool. +func NewTransferToAgentTool(curAgent, parent agent.Agent, targets []agent.Agent) (*TransferToAgentTool, error) { + si, err := instructionsForTransferToAgent(curAgent, parent, targets) + if err != nil { + return nil, err + } + return &TransferToAgentTool{ + instructions: si, + supportedAgents: targets, + }, nil +} // Description implements tool.Tool. func (t *TransferToAgentTool) Description() string { @@ -106,7 +124,7 @@ This tool hands off control to another agent when it's more suitable to answer t // Name implements tool.Tool. func (t *TransferToAgentTool) Name() string { - return "transfer_to_agent" + return transferAgentName } // IsLongRunning implements tool.Tool. @@ -119,11 +137,12 @@ func (t *TransferToAgentTool) Declaration() *genai.FunctionDeclaration { Name: t.Name(), Description: t.Description(), Parameters: &genai.Schema{ - Type: "object", + Type: genai.TypeObject, Properties: map[string]*genai.Schema{ "agent_name": { - Type: "string", + Type: genai.TypeString, Description: "the agent name to transfer to", + Enum: t.enums(), }, }, Required: []string{"agent_name"}, @@ -131,13 +150,21 @@ func (t *TransferToAgentTool) Declaration() *genai.FunctionDeclaration { } } +func (t *TransferToAgentTool) enums() []string { + var agentNames []string + for _, a := range t.supportedAgents { + agentNames = append(agentNames, a.Name()) + } + return agentNames +} + // ProcessRequest implements types.Tool. -func (t *TransferToAgentTool) ProcessRequest(ctx tool.Context, req *model.LLMRequest) error { +func (t *TransferToAgentTool) ProcessRequest(ctx agent.ToolContext, req *model.LLMRequest) error { return appendTools(req, t) } // Run implements types.Tool. -func (t *TransferToAgentTool) Run(ctx tool.Context, args any) (map[string]any, error) { +func (t *TransferToAgentTool) Run(ctx agent.ToolContext, args any) (map[string]any, error) { if args == nil { return nil, fmt.Errorf("missing argument") } @@ -201,7 +228,7 @@ func shouldUseAutoFlow(agent agent.Agent) bool { return len(agent.SubAgents()) != 0 || !a.internal().DisallowTransferToParent || !a.internal().DisallowTransferToPeers } -// AppendTools appends the tools to the request. +// appendTools appends the tools to the request. // Appending duplicate tools or nameless tools is an error. func appendTools(r *model.LLMRequest, tools ...tool.Tool) error { if r.Tools == nil { @@ -254,7 +281,7 @@ func appendTools(r *model.LLMRequest, tools ...tool.Tool) error { var transferToAgentPromptTmpl = template.Must( template.New("transfer_to_agent_prompt").Parse(agentTransferInstructionTemplate)) -func instructionsForTransferToAgent(curAgent, parent agent.Agent, targets []agent.Agent, transferTool tool.Tool) (string, error) { +func instructionsForTransferToAgent(curAgent, parent agent.Agent, targets []agent.Agent) (string, error) { if asLLMAgent(curAgent).internal().DisallowTransferToParent { parent = nil } @@ -270,7 +297,7 @@ func instructionsForTransferToAgent(curAgent, parent agent.Agent, targets []agen AgentName: curAgent.Name(), Parent: parent, Targets: targets, - ToolName: transferTool.Name(), + ToolName: transferAgentName, FormattedTargets: formatTargets(targets), }); err != nil { return "", err diff --git a/internal/llminternal/agent_transfer_test.go b/internal/llminternal/agent_transfer_test.go index 356e1e906..67ae41099 100644 --- a/internal/llminternal/agent_transfer_test.go +++ b/internal/llminternal/agent_transfer_test.go @@ -21,6 +21,7 @@ import ( "testing" "github.com/google/go-cmp/cmp" + "github.com/google/go-cmp/cmp/cmpopts" "google.golang.org/genai" "google.golang.org/adk/agent" @@ -47,6 +48,7 @@ func TestAgentTransferRequestProcessor(t *testing.T) { } check := func(t *testing.T, curAgent, root agent.Agent, wantParent string, wantAgents, unwantAgents []string) { + t.Helper() req := &model.LLMRequest{} parents, err := parentmap.New(root) @@ -118,6 +120,22 @@ func TestAgentTransferRequestProcessor(t *testing.T) { }) { t.Errorf("instruction does not include subagents, got: %s", strings.Join(instructions, "\n")) } + if len(wantAgents) > 0 { + transferTool, ok := gotTool.(*llminternal.TransferToAgentTool) + if !ok { + t.Fatalf("failed to type convert tool %v, got %T", wantToolName, gotTool) + } + declaration := transferTool.Declaration() + gotEnums := slices.Clone(declaration.Parameters.Properties["agent_name"].Enum) + wantEnums := slices.Clone(wantAgents) + // Add parent to the list of agents to transfer to, if not already present. + if wantParent != "" && !slices.Contains(wantAgents, wantParent) { + wantEnums = append(wantEnums, wantParent) + } + if diff := cmp.Diff(wantEnums, gotEnums, cmpopts.SortSlices(func(a, b string) bool { return a < b })); diff != "" { + t.Fatalf("failed to set agent_name enums (-want, +got): %v", diff) + } + } if len(unwantAgents) > 0 && slices.ContainsFunc(instructions, func(s string) bool { return slices.ContainsFunc(unwantAgents, func(unwanted string) bool { for _, unwanted := range unwantAgents { @@ -417,7 +435,7 @@ func TestAgentTransfer_ProcessRequest(t *testing.T) { x int } var req model.LLMRequest - handler := func(ctx tool.Context, input Input) (int, error) { + handler := func(ctx agent.ToolContext, input Input) (int, error) { return input.x, nil } identityTool, err := functiontool.New(functiontool.Config{ @@ -453,7 +471,7 @@ func TestTransferToAgentToolRun(t *testing.T) { curTool := &llminternal.TransferToAgentTool{} invCtx := icontext.NewInvocationContext(t.Context(), icontext.InvocationContextParams{}) - ctx := toolinternal.NewToolContext(invCtx, "", &session.EventActions{}, nil) + ctx := agent.NewToolContext(invCtx, "", &session.EventActions{}, nil) wantAgentName := "TestAgent" args := map[string]any{"agent_name": wantAgentName} @@ -479,7 +497,7 @@ func TestTransferToAgentToolRun(t *testing.T) { for _, tc := range testCases { t.Run(tc.name, func(t *testing.T) { curTool := &llminternal.TransferToAgentTool{} - ctx := toolinternal.NewToolContext(icontext.NewInvocationContext(t.Context(), icontext.InvocationContextParams{}), "", nil, nil) + ctx := agent.NewToolContext(icontext.NewInvocationContext(t.Context(), icontext.InvocationContextParams{}), "", nil, nil) if got, err := curTool.Run(ctx, tc.args); err == nil { t.Fatalf("Run(%v) = (%v, %v), want error", tc.args, got, err) } diff --git a/internal/llminternal/audio_cache_manager.go b/internal/llminternal/audio_cache_manager.go new file mode 100644 index 000000000..4446ddbbb --- /dev/null +++ b/internal/llminternal/audio_cache_manager.go @@ -0,0 +1,171 @@ +// Copyright 2026 Google LLC +// +// Licensed under the Apache License, Version 2.0 (the "License"); +// you may not use this file except in compliance with the License. +// You may obtain a copy of the License at +// +// http://www.apache.org/licenses/LICENSE-2.0 +// +// Unless required by applicable law or agreed to in writing, software +// distributed under the License is distributed on an "AS IS" BASIS, +// WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. +// See the License for the specific language governing permissions and +// limitations under the License. + +package llminternal + +import ( + "bytes" + "context" + "fmt" + "strings" + "sync" + "time" + + "google.golang.org/genai" + + "google.golang.org/adk/agent" + "google.golang.org/adk/platform" + "google.golang.org/adk/session" +) + +// AudioCacheManager manages audio caching and flushing for live streaming flows. +type AudioCacheManager struct { + mu sync.Mutex + inputCache [][]byte + outputCache [][]byte + inputStartTime time.Time + outputStartTime time.Time + inputMimeType string + outputMimeType string +} + +// NewAudioCacheManager creates a new AudioCacheManager. +func NewAudioCacheManager() *AudioCacheManager { + return &AudioCacheManager{ + inputMimeType: "audio/pcm", // Default to audio/pcm + outputMimeType: "audio/pcm", // Default to audio/pcm + } +} + +// CacheInput caches incoming user audio data. +func (m *AudioCacheManager) CacheInput(ctx context.Context, data []byte, mimeType string) { + if mimeType != "" && !strings.HasPrefix(mimeType, "audio/") { + return + } + m.mu.Lock() + defer m.mu.Unlock() + if len(m.inputCache) == 0 { + m.inputStartTime = platform.Now(ctx) + if mimeType != "" { + m.inputMimeType = mimeType + } + } + m.inputCache = append(m.inputCache, data) +} + +// CacheOutput caches outgoing model audio data. +func (m *AudioCacheManager) CacheOutput(ctx context.Context, data []byte, mimeType string) { + if mimeType != "" && !strings.HasPrefix(mimeType, "audio/") { + return + } + m.mu.Lock() + defer m.mu.Unlock() + if len(m.outputCache) == 0 { + m.outputStartTime = platform.Now(ctx) + if mimeType != "" { + m.outputMimeType = mimeType + } + } + m.outputCache = append(m.outputCache, data) +} + +// FlushCaches flushes audio caches to artifact services and returns created events. +func (m *AudioCacheManager) FlushCaches(ctx agent.InvocationContext, flushUser, flushModel bool) ([]*session.Event, error) { + m.mu.Lock() + defer m.mu.Unlock() + + var events []*session.Event + + if flushUser && len(m.inputCache) > 0 { + ev, err := m.flushCache(ctx, m.inputCache, "input_audio", m.inputMimeType, m.inputStartTime) + if err != nil { + return nil, err + } + if ev != nil { + events = append(events, ev) + } + m.inputCache = nil + } + + if flushModel && len(m.outputCache) > 0 { + ev, err := m.flushCache(ctx, m.outputCache, "output_audio", m.outputMimeType, m.outputStartTime) + if err != nil { + return nil, err + } + if ev != nil { + events = append(events, ev) + } + m.outputCache = nil + } + + return events, nil +} + +func (m *AudioCacheManager) flushCache(ctx agent.InvocationContext, cache [][]byte, cacheType, mimeType string, startTime time.Time) (*session.Event, error) { + if len(cache) == 0 { + return nil, nil + } + + // Combine chunks + var buf bytes.Buffer + for _, chunk := range cache { + buf.Write(chunk) + } + combinedData := buf.Bytes() + + // Generate filename + timestamp := startTime.UnixMilli() + ext := "pcm" + if mimeType == "audio/mp3" { + ext = "mp3" + } + + filename := fmt.Sprintf("adk_live_audio_storage_%s_%d.%s", cacheType, timestamp, ext) + + // Save to artifact service + part := genai.NewPartFromBytes(combinedData, mimeType) + resp, err := ctx.Artifacts().Save(ctx, filename, part) + if err != nil { + return nil, fmt.Errorf("failed to save audio artifact: %w", err) + } + + // Create artifact reference + sess := ctx.Session() + artifactRef := fmt.Sprintf("artifact://%s/%s/%s/_adk_live/%s#%d", sess.AppName(), sess.UserID(), sess.ID(), filename, resp.Version) + + // Create event with file data reference + author := ctx.Agent().Name() + role := "model" + if cacheType == "input_audio" { + author = "user" + role = "user" + } + + ev := session.NewEventWithContext(ctx, ctx.InvocationID()) + ev.Author = author + ev.Timestamp = startTime + ev.Content = &genai.Content{ + Role: role, + Parts: []*genai.Part{ + { + FileData: &genai.FileData{ + FileURI: artifactRef, + MIMEType: mimeType, + }, + }, + }, + } + + return ev, nil +} diff --git a/internal/llminternal/audio_cache_manager_test.go b/internal/llminternal/audio_cache_manager_test.go new file mode 100644 index 000000000..91b7cc448 --- /dev/null +++ b/internal/llminternal/audio_cache_manager_test.go @@ -0,0 +1,289 @@ +// Copyright 2026 Google LLC +// +// Licensed under the Apache License, Version 2.0 (the "License"); +// you may not use this file except in compliance with the License. +// You may obtain a copy of the License at +// +// http://www.apache.org/licenses/LICENSE-2.0 +// +// Unless required by applicable law or agreed to in writing, software +// distributed under the License is distributed on an "AS IS" BASIS, +// WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. +// See the License for the specific language governing permissions and +// limitations under the License. + +package llminternal + +import ( + "bytes" + "context" + "testing" + "time" + + "google.golang.org/genai" + + "google.golang.org/adk/agent" + "google.golang.org/adk/artifact" + "google.golang.org/adk/session" +) + +type audioMockArtifacts struct { + savedName string + savedPart *genai.Part +} + +func (m *audioMockArtifacts) Save(ctx context.Context, name string, data *genai.Part) (*artifact.SaveResponse, error) { + m.savedName = name + m.savedPart = data + return &artifact.SaveResponse{Version: 1}, nil +} + +func (m *audioMockArtifacts) List(context.Context) (*artifact.ListResponse, error) { return nil, nil } +func (m *audioMockArtifacts) Load(ctx context.Context, name string) (*artifact.LoadResponse, error) { + return nil, nil +} + +func (m *audioMockArtifacts) LoadVersion(ctx context.Context, name string, version int) (*artifact.LoadResponse, error) { + return nil, nil +} + +type audioMockSession struct { + session.Session + id string + appName string + userID string +} + +func (m *audioMockSession) ID() string { return m.id } +func (m *audioMockSession) AppName() string { return m.appName } +func (m *audioMockSession) UserID() string { return m.userID } + +type audioMockInvocationContext struct { + agent.InvocationContext + artifacts agent.Artifacts + session session.Session + invocationID string + agentObj agent.Agent +} + +func (m *audioMockInvocationContext) Artifacts() agent.Artifacts { return m.artifacts } +func (m *audioMockInvocationContext) Session() session.Session { return m.session } +func (m *audioMockInvocationContext) InvocationID() string { return m.invocationID } +func (m *audioMockInvocationContext) Agent() agent.Agent { return m.agentObj } + +func (m *audioMockInvocationContext) Deadline() (time.Time, bool) { return time.Time{}, false } +func (m *audioMockInvocationContext) Done() <-chan struct{} { return nil } +func (m *audioMockInvocationContext) Err() error { return nil } +func (m *audioMockInvocationContext) Value(any) any { return nil } + +type audioMockAgent struct { + agent.Agent + name string +} + +func (m *audioMockAgent) Name() string { return m.name } + +func TestAudioCacheManager(t *testing.T) { + type chunk struct { + data []byte + mime string + } + tests := []struct { + name string + inputs []chunk + outputs []chunk + flushUser bool + flushModel bool + expectedEventsCount int + verify func(t *testing.T, events []*session.Event, mockArt *audioMockArtifacts) + }{ + { + name: "FlushBoth_PCM", + inputs: []chunk{ + {[]byte("input1"), "audio/pcm"}, + {[]byte("input2"), "audio/pcm"}, + }, + outputs: []chunk{ + {[]byte("output1"), "audio/pcm"}, + {[]byte("output2"), "audio/pcm"}, + }, + flushUser: true, + flushModel: true, + expectedEventsCount: 2, + verify: func(t *testing.T, events []*session.Event, mockArt *audioMockArtifacts) { + // Verify input event + ev1 := events[0] + if ev1.Author != "user" || ev1.Content.Role != "user" { + t.Errorf("ev1 author/role mismatch: author=%q, role=%q", ev1.Author, ev1.Content.Role) + } + if ev1.Content.Parts[0].FileData.MIMEType != "audio/pcm" { + t.Errorf("ev1 mimeType mismatch: got %s", ev1.Content.Parts[0].FileData.MIMEType) + } + + // Verify output event + ev2 := events[1] + if ev2.Author != "agent1" || ev2.Content.Role != "model" { + t.Errorf("ev2 author/role mismatch: author=%q, role=%q", ev2.Author, ev2.Content.Role) + } + if ev2.Content.Parts[0].FileData.MIMEType != "audio/pcm" { + t.Errorf("ev2 mimeType mismatch: got %s", ev2.Content.Parts[0].FileData.MIMEType) + } + }, + }, + { + name: "FlushSelective_InputOnly", + inputs: []chunk{ + {[]byte("input1"), "audio/pcm"}, + }, + outputs: []chunk{ + {[]byte("output1"), "audio/pcm"}, + }, + flushUser: true, + flushModel: false, + expectedEventsCount: 1, + verify: func(t *testing.T, events []*session.Event, mockArt *audioMockArtifacts) { + if events[0].Author != "user" { + t.Errorf("Expected author user, got %s", events[0].Author) + } + }, + }, + { + name: "FlushSelective_OutputOnly", + inputs: []chunk{ + {[]byte("input1"), "audio/pcm"}, + }, + outputs: []chunk{ + {[]byte("output1"), "audio/pcm"}, + }, + flushUser: false, + flushModel: true, + expectedEventsCount: 1, + verify: func(t *testing.T, events []*session.Event, mockArt *audioMockArtifacts) { + if events[0].Author != "agent1" { + t.Errorf("Expected author agent1, got %s", events[0].Author) + } + }, + }, + { + name: "FlushEmpty", + flushUser: true, + flushModel: true, + expectedEventsCount: 0, + }, + { + name: "VerifyCombinedData", + inputs: []chunk{ + {[]byte("chunk1"), "audio/pcm"}, + {[]byte("chunk2"), "audio/pcm"}, + }, + flushUser: true, + flushModel: false, + expectedEventsCount: 1, + verify: func(t *testing.T, events []*session.Event, mockArt *audioMockArtifacts) { + if mockArt.savedPart == nil { + t.Fatal("Expected savedPart, got nil") + } + expectedData := []byte("chunk1chunk2") + if !bytes.Equal(mockArt.savedPart.InlineData.Data, expectedData) { + t.Errorf("Expected combined data %s, got %s", expectedData, mockArt.savedPart.InlineData.Data) + } + }, + }, + { + name: "MimeTypeFallback", + inputs: []chunk{ + {[]byte("input1"), ""}, + }, + flushUser: true, + flushModel: false, + expectedEventsCount: 1, + verify: func(t *testing.T, events []*session.Event, mockArt *audioMockArtifacts) { + if mockArt.savedPart.InlineData.MIMEType != "audio/pcm" { + t.Errorf("Expected fallback MIMEType audio/pcm, got %s", mockArt.savedPart.InlineData.MIMEType) + } + }, + }, + { + name: "DifferentMimeTypes", + inputs: []chunk{ + {[]byte("input1"), "audio/pcm"}, + }, + outputs: []chunk{ + {[]byte("output1"), "audio/mp3"}, + }, + flushUser: true, + flushModel: true, + expectedEventsCount: 2, + verify: func(t *testing.T, events []*session.Event, mockArt *audioMockArtifacts) { + if events[0].Content.Parts[0].FileData.MIMEType != "audio/pcm" { + t.Errorf("Expected input MIMEType audio/pcm, got %s", events[0].Content.Parts[0].FileData.MIMEType) + } + if events[1].Content.Parts[0].FileData.MIMEType != "audio/mp3" { + t.Errorf("Expected output MIMEType audio/mp3, got %s", events[1].Content.Parts[0].FileData.MIMEType) + } + }, + }, + { + name: "FiltersNonAudio", + inputs: []chunk{ + {[]byte("input1"), "video/mp4"}, + {[]byte("input2"), "image/png"}, + {[]byte("audio_input"), "audio/pcm"}, + }, + outputs: []chunk{ + {[]byte("output1"), "video/h264"}, + }, + flushUser: true, + flushModel: true, + expectedEventsCount: 1, + verify: func(t *testing.T, events []*session.Event, mockArt *audioMockArtifacts) { + ev := events[0] + if ev.Author != "user" { + t.Errorf("Expected author user, got %s", ev.Author) + } + if mockArt.savedPart == nil { + t.Fatal("Expected savedPart, got nil") + } + if !bytes.Equal(mockArt.savedPart.InlineData.Data, []byte("audio_input")) { + t.Errorf("Expected only 'audio_input' to be saved, got %s", mockArt.savedPart.InlineData.Data) + } + }, + }, + } + + for _, tt := range tests { + t.Run(tt.name, func(t *testing.T) { + mgr := NewAudioCacheManager() + + for _, in := range tt.inputs { + mgr.CacheInput(context.Background(), in.data, in.mime) + } + for _, out := range tt.outputs { + mgr.CacheOutput(context.Background(), out.data, out.mime) + } + + mockArt := &audioMockArtifacts{} + mockSess := &audioMockSession{id: "sess1", appName: "app1", userID: "user1"} + mockAg := &audioMockAgent{name: "agent1"} + mockCtx := &audioMockInvocationContext{ + artifacts: mockArt, + session: mockSess, + invocationID: "inv1", + agentObj: mockAg, + } + + events, err := mgr.FlushCaches(mockCtx, tt.flushUser, tt.flushModel) + if err != nil { + t.Fatalf("FlushCaches failed: %v", err) + } + + if len(events) != tt.expectedEventsCount { + t.Fatalf("Expected %d events, got %d", tt.expectedEventsCount, len(events)) + } + + if tt.verify != nil { + tt.verify(t, events, mockArt) + } + }) + } +} diff --git a/internal/llminternal/base_flow.go b/internal/llminternal/base_flow.go index 2934fd165..8b0839d6f 100644 --- a/internal/llminternal/base_flow.go +++ b/internal/llminternal/base_flow.go @@ -18,13 +18,15 @@ import ( "context" "errors" "fmt" + "io" "iter" + "log" "maps" "slices" "strings" "sync" + "time" - "github.com/google/uuid" "google.golang.org/genai" "google.golang.org/adk/agent" @@ -38,29 +40,43 @@ import ( "google.golang.org/adk/internal/toolinternal" "google.golang.org/adk/internal/utils" "google.golang.org/adk/model" + "google.golang.org/adk/platform" "google.golang.org/adk/session" "google.golang.org/adk/tool" "google.golang.org/adk/tool/toolconfirmation" ) +// ErrModelNotConfigured is returned when the model is not configured. var ErrModelNotConfigured = errors.New("model not configured; ensure Model is set in llmagent.Config") +// BeforeModelCallback is called before sending a request to the model. +// +// If it returns non-nil result or error, the actual call is skipped and the returned value is used +// as the model invocation result. type BeforeModelCallback func(ctx agent.CallbackContext, llmRequest *model.LLMRequest) (*model.LLMResponse, error) +// AfterModelCallback is called after receiving a response from the model. +// +// If it returns non-nil result or error, the returned value is used as the model invocation result. type AfterModelCallback func(ctx agent.CallbackContext, llmResponse *model.LLMResponse, llmResponseError error) (*model.LLMResponse, error) +// OnModelErrorCallback is called when the model returns an error. +// +// If it returns non-nil result or error, the returned value is used as the model invocation result. type OnModelErrorCallback func(ctx agent.CallbackContext, llmRequest *model.LLMRequest, llmResponseError error) (*model.LLMResponse, error) // Tool-callback types live in the liveflow subpackage so it can stand on its // own without importing back into llminternal. The non-Live flow continues to // reference them via these aliases for backward compatibility — both names -// refer to the exact same underlying types. +// refer to the exact same underlying types. Signatures match upstream's +// (agent.ToolContext == tool.Context via alias). type ( BeforeToolCallback = liveflow.BeforeToolCallback AfterToolCallback = liveflow.AfterToolCallback OnToolErrorCallback = liveflow.OnToolErrorCallback ) +// Flow is a base implementation of the agent flow. type Flow struct { Model model.LLM @@ -128,6 +144,405 @@ func (f *Flow) Run(ctx agent.InvocationContext) iter.Seq2[*session.Event, error] } } +type activeTask struct { + callID string + cancel context.CancelFunc +} + +type liveSessionImpl struct { + inputCh chan agent.LiveRequest + outputCh chan eventOrError + done chan struct{} + closeOnce sync.Once + audioMgr *AudioCacheManager + mu sync.Mutex + activeTools map[string][]activeTask +} + +func (s *liveSessionImpl) RegisterStreamingTool(toolName, callID string, cancel context.CancelFunc) { + s.mu.Lock() + defer s.mu.Unlock() + if s.activeTools == nil { + s.activeTools = make(map[string][]activeTask) + } + s.activeTools[toolName] = append(s.activeTools[toolName], activeTask{ + callID: callID, + cancel: cancel, + }) +} + +func (s *liveSessionImpl) UnregisterStreamingTool(toolName, callID string) { + s.mu.Lock() + defer s.mu.Unlock() + tasks, exists := s.activeTools[toolName] + if !exists { + return + } + for i, task := range tasks { + if task.callID == callID { + s.activeTools[toolName] = append(tasks[:i], tasks[i+1:]...) + break + } + } + if len(s.activeTools[toolName]) == 0 { + delete(s.activeTools, toolName) + } +} + +func (s *liveSessionImpl) CancelAllStreamingTools(toolName string) bool { + s.mu.Lock() + defer s.mu.Unlock() + tasks, exists := s.activeTools[toolName] + if !exists || len(tasks) == 0 { + return false + } + for _, task := range tasks { + task.cancel() + } + delete(s.activeTools, toolName) + return true +} + +type eventOrError struct { + event *session.Event + err error +} + +func newLiveSessionImpl() *liveSessionImpl { + return &liveSessionImpl{ + inputCh: make(chan agent.LiveRequest), + outputCh: make(chan eventOrError), + done: make(chan struct{}), + audioMgr: NewAudioCacheManager(), + } +} + +func (s *liveSessionImpl) Send(req agent.LiveRequest) error { + select { + case s.inputCh <- req: + return nil + case <-s.done: + return io.EOF + } +} + +func (s *liveSessionImpl) recvIter() iter.Seq2[*session.Event, error] { + return func(yield func(*session.Event, error) bool) { + for { + select { + case res := <-s.outputCh: + if !yield(res.event, res.err) { + return + } + case <-s.done: + return + } + } + } +} + +func (s *liveSessionImpl) Close() error { + s.closeOnce.Do(func() { + close(s.done) + }) + return nil +} + +func (s *liveSessionImpl) pushEvent(ev *session.Event) bool { + select { + case s.outputCh <- eventOrError{event: ev}: + return true + case <-s.done: + return false + } +} + +func (s *liveSessionImpl) pushError(err error) bool { + select { + case s.outputCh <- eventOrError{err: err}: + return true + case <-s.done: + return false + } +} + +func (f *Flow) RunLive(ctx agent.InvocationContext) (agent.LiveSession, iter.Seq2[*session.Event, error], error) { + clientProvider, ok := f.Model.(interface { + Client() *genai.Client + }) + if !ok { + return nil, nil, fmt.Errorf("model does not support live connection") + } + client := clientProvider.Client() + + runCfg := runconfig.FromContext(ctx) + if runCfg == nil || runCfg.Live == nil { + return nil, nil, fmt.Errorf("live run config not found") + } + + sess := newLiveSessionImpl() + + go func() { + defer func() { + _ = sess.Close() + }() + + nreq := &model.LLMRequest{ + Model: f.Model.Name(), + } + for ev, err := range f.preprocess(ctx, nreq) { + if err != nil { + sess.pushError(err) + return + } + if ev != nil { + if !sess.pushEvent(ev) { + return + } + } + } + + liveConnectConfig := &genai.LiveConnectConfig{ + ResponseModalities: runCfg.Live.ResponseModalities, + SpeechConfig: runCfg.Live.SpeechConfig, + SystemInstruction: nreq.Config.SystemInstruction, + Tools: nreq.Config.Tools, + SessionResumption: runCfg.Live.SessionResumption, + InputAudioTranscription: runCfg.Live.InputAudioTranscription, + OutputAudioTranscription: runCfg.Live.OutputAudioTranscription, + } + + isResumable := func(err error) bool { + if err == nil { + return false + } + if err == io.EOF { + return true + } + errStr := err.Error() + return strings.Contains(errStr, "broken pipe") || + strings.Contains(errStr, "connection reset") || + strings.Contains(errStr, "EOF") || + strings.Contains(errStr, "1008") || + strings.Contains(errStr, "GoAway") + } + + iCtx, isIContext := ctx.(*icontext.InvocationContext) + + for { + if isIContext { + handle := iCtx.LiveSessionResumptionHandle() + if handle != "" { + if liveConnectConfig.SessionResumption == nil { + liveConnectConfig.SessionResumption = &genai.SessionResumptionConfig{} + } + liveConnectConfig.SessionResumption.Handle = handle + if googlellm.GetGoogleLLMVariant(f.Model) == genai.BackendVertexAI { + liveConnectConfig.SessionResumption.Transparent = true + } + } + } + + connCtx, cancelConn := context.WithCancel(ctx) + + if liveConnectConfig.SessionResumption != nil { + log.Printf("connecting with live session handle: %s\n", liveConnectConfig.SessionResumption.Handle) + } + liveSession, err := client.Live.Connect(connCtx, f.Model.Name(), liveConnectConfig) + if err != nil { + cancelConn() + log.Printf("failed to connect live session: %v\n", err) + sess.pushError(fmt.Errorf("failed to connect live session: %w", err)) + return + } + + liveConn := googlellm.NewLiveConnection(liveSession, f.Model.Name(), googlellm.GetGoogleLLMVariant(f.Model)) + + cleanup := func() { + cancelConn() + _ = liveConn.Close() + } + + eventsChan := make(chan *session.Event) + errChan := make(chan error) + + // Send preprocessed content directly to model if any exists after early preprocessing + if len(nreq.Contents) > 0 { + if err := liveConn.SendHistory(ctx, nreq.Contents); err != nil { + log.Printf("failed to send history: %v\n", err) + sess.pushError(err) + return + } + } + + // Reading from model loop + go func() { + for { + resp, err := liveConn.Recv(connCtx) + if err != nil { + errChan <- err + return + } + if resp != nil { + if resp.SessionResumptionHandle != "" { + if isIContext { + log.Printf("received session resumption handle: %s\n", resp.SessionResumptionHandle) + iCtx.SetLiveSessionResumptionHandle(resp.SessionResumptionHandle) + } + } + if runCfg.Live.SaveLiveBlob && resp.Content != nil { + for _, part := range resp.Content.Parts { + if part.InlineData != nil { + sess.audioMgr.CacheOutput(ctx, part.InlineData.Data, part.InlineData.MIMEType) + } + } + } + ev := session.NewEventWithContext(ctx, ctx.InvocationID()) + ev.Author = ctx.Agent().Name() + ev.LLMResponse = *resp + select { + case eventsChan <- ev: + case <-connCtx.Done(): + return + } + } + } + }() + + // Sending to model loop + go func() { + for { + select { + case <-connCtx.Done(): + return + case req, ok := <-sess.inputCh: + if !ok { + return + } + if req.Content != nil { + if err := liveConn.SendContent(connCtx, req.Content); err != nil { + errChan <- err + return + } + } + if req.RealtimeInput != nil { + if blob, ok := req.RealtimeInput.(*genai.Blob); ok { + sess.audioMgr.CacheInput(ctx, blob.Data, blob.MIMEType) + } + if err := liveConn.SendRealtime(connCtx, req.RealtimeInput); err != nil { + errChan <- err + return + } + } + } + } + }() + + reconnect := false + for !reconnect { + select { + case ev := <-eventsChan: + if !sess.pushEvent(ev) { + cleanup() + return + } + // Flush caches if needed + if runCfg.Live.SaveLiveBlob { + var flushUser, flushModel bool + if ev.LLMResponse.Interrupted { + flushModel = true + } + if ev.LLMResponse.TurnComplete { + flushUser = true + flushModel = true + } + if flushUser || flushModel { + flushedEvents, err := sess.audioMgr.FlushCaches(ctx, flushUser, flushModel) + if err != nil { + sess.pushError(err) + cleanup() + return + } + for _, fev := range flushedEvents { + if !sess.pushEvent(fev) { + cleanup() + return + } + } + } + } + // Handle function calls if present in the event + fnCalls := utils.FunctionCalls(ev.LLMResponse.Content) + if len(fnCalls) > 0 { + tools := make(map[string]tool.Tool) + for _, t := range f.Tools { + tools[t.Name()] = t + } + respEv, err := f.handleFunctionCalls(ctx, tools, &ev.LLMResponse, nil, sess) + if err != nil { + sess.pushError(err) + cleanup() + return + } + if respEv != nil { + if !sess.pushEvent(respEv) { + cleanup() + return + } + // Check if task_completed was invoked. + var isTaskCompleted bool + if respEv.LLMResponse.Content != nil { + for _, part := range respEv.LLMResponse.Content.Parts { + if part.FunctionResponse != nil && part.FunctionResponse.Name == "task_completed" { + isTaskCompleted = true + break + } + } + } + if isTaskCompleted { + time.Sleep(100 * time.Millisecond) + cleanup() + return + } + // Send function response back to model + if err := liveConn.SendContent(connCtx, respEv.LLMResponse.Content); err != nil { + sess.pushError(err) + cleanup() + return + } + } + } + case err := <-errChan: + if isResumable(err) { + log.Printf("Connection error, attempting to resume: %v\n", err) + reconnect = true + break // Break the select + } + sess.pushError(err) + cleanup() + return + case <-ctx.Done(): + sess.pushError(ctx.Err()) + cleanup() + return + } + if reconnect { + break // Break the for loop + } + } + + // Cleanup before reconnecting + cleanup() + + if !reconnect { + break + } + } + }() + + return sess, sess.recvIter(), nil +} + func (f *Flow) runOneStep(ctx agent.InvocationContext) iter.Seq2[*session.Event, error] { return func(yield func(*session.Event, error) bool) { if f.Model == nil { @@ -182,6 +597,7 @@ func (f *Flow) runOneStep(ctx agent.InvocationContext) iter.Seq2[*session.Event, if !yield(nil, fmt.Errorf("unexpected tool type %T for tool %v", v, k)) { return } + continue } tools[k] = tool } @@ -198,7 +614,7 @@ func (f *Flow) runOneStep(ctx agent.InvocationContext) iter.Seq2[*session.Event, } // Handle function calls. - ev, err := f.handleFunctionCalls(ctx, tools, resp.LLMResponse, nil) + ev, err := f.handleFunctionCalls(ctx, tools, resp.LLMResponse, nil, nil) if err != nil { yield(nil, err) return @@ -209,16 +625,19 @@ func (f *Flow) runOneStep(ctx agent.InvocationContext) iter.Seq2[*session.Event, } toolConfirmationEvent := generateRequestConfirmationEvent(ctx, modelResponseEvent, ev) + + // Yield function responses before confirmation requests so consumers that + // pause for user approval still persist completed tool results. + if !yield(ev, nil) { + return + } + if toolConfirmationEvent != nil { if !yield(toolConfirmationEvent, nil) { return } } - if !yield(ev, nil) { - return - } - // If the model response is structured, yield it as a final model response event. outputSchemaResponse, err := retrieveStructuredModelResponse(ev) if err != nil { @@ -288,7 +707,7 @@ func toolPreprocess(ctx agent.InvocationContext, req *model.LLMRequest, tools [] return fmt.Errorf("tool %q does not implement RequestProcessor() method", t.Name()) } // TODO: how to prevent mutation on this? - toolCtx := toolinternal.NewToolContext(ctx, "", &session.EventActions{}, nil) + toolCtx := agent.NewToolContext(ctx, "", &session.EventActions{}, nil) if err := requestProcessor.ProcessRequest(toolCtx, req); err != nil { return err } @@ -306,7 +725,7 @@ func toolsetPreprocess(ctx agent.InvocationContext, req *model.LLMRequest) error if !ok { continue // Not all toolsets implement RequestProcessor. } - toolCtx := toolinternal.NewToolContext(ctx, "", nil, nil) + toolCtx := agent.NewToolContext(ctx, "", nil, nil) if err := processor.ProcessRequest(toolCtx, req); err != nil { return fmt.Errorf("process request by toolset %q: %w", toolset.Name(), err) } @@ -314,8 +733,8 @@ func toolsetPreprocess(ctx agent.InvocationContext, req *model.LLMRequest) error return nil } -func newResponseWithEventID(resp *model.LLMResponse) *responseWithEventID { - return &responseWithEventID{resp, uuid.New().String()} +func newResponseWithEventID(ctx context.Context, resp *model.LLMResponse) *responseWithEventID { + return &responseWithEventID{resp, platform.NewUUID(ctx)} } func (f *Flow) callLLM(ctx agent.InvocationContext, req *model.LLMRequest, stateDelta map[string]any, artifactDelta map[string]int64) iter.Seq2[*responseWithEventID, error] { @@ -325,7 +744,7 @@ func (f *Flow) callLLM(ctx agent.InvocationContext, req *model.LLMRequest, state cctx := icontext.NewCallbackContextWithDelta(ctx, stateDelta, artifactDelta) callbackResponse, callbackErr := pluginManager.RunBeforeModelCallback(cctx, req) if callbackResponse != nil || callbackErr != nil { - yield(newResponseWithEventID(callbackResponse), callbackErr) + yield(newResponseWithEventID(ctx, callbackResponse), callbackErr) return } } @@ -335,7 +754,7 @@ func (f *Flow) callLLM(ctx agent.InvocationContext, req *model.LLMRequest, state callbackResponse, callbackErr := callback(cctx, req) if callbackResponse != nil || callbackErr != nil { - yield(newResponseWithEventID(callbackResponse), callbackErr) + yield(newResponseWithEventID(ctx, callbackResponse), callbackErr) return } } @@ -365,7 +784,7 @@ func (f *Flow) callLLM(ctx agent.InvocationContext, req *model.LLMRequest, state } // Function call ID is optional in genai API and some models do not use the field. // Set it in case after model callbacks use it. - utils.PopulateClientFunctionCallID(resp.Content) + utils.PopulateClientFunctionCallID(ctx, resp.Content) callbackResp, callbackErr := f.runAfterModelCallbacks(ctx, resp.LLMResponse, stateDelta, artifactDelta, err) // TODO: check if we should stop iterator on the first error from stream or continue yielding next results. @@ -435,7 +854,7 @@ func generateContent(ctx agent.InvocationContext, m model.LLM, req *model.LLMReq // Ensure that the span is ended in case of error or if none final responses are yielded before the yield returns false. defer endSpanAndTrackResult() for resp, err := range m.GenerateContent(ctx, req, useStream) { - response := newResponseWithEventID(resp) + response := newResponseWithEventID(ctx, resp) lastResponse = *response lastErr = err // Complete the span immediately to avoid capturing the upstream yield processing time. @@ -525,9 +944,9 @@ func (f *Flow) finalizeModelResponseEvent(ctx agent.InvocationContext, resp *res // FunctionCall & FunctionResponse matching algorithm assumes non-empty function call IDs // but function call ID is optional in genai API and some models do not use the field. // Generate function call ids. (see functions.populate_client_function_call_id in python SDK) - utils.PopulateClientFunctionCallID(resp.Content) + utils.PopulateClientFunctionCallID(ctx, resp.Content) - ev := session.NewEvent(ctx.InvocationID()) + ev := session.NewEventWithContext(ctx, ctx.InvocationID()) ev.ID = resp.eventID // TODO change NewEvent to accept event id ev.Author = ctx.Agent().Name() ev.Branch = ctx.Branch() @@ -583,11 +1002,32 @@ Suggested fixes: - Check for typos in function name`, toolName, joinedTools) } +type cancelledToolContext struct { + agent.ToolContext + cancelCtx context.Context +} + +func (c *cancelledToolContext) Done() <-chan struct{} { + return c.cancelCtx.Done() +} + +func (c *cancelledToolContext) Err() error { + return c.cancelCtx.Err() +} + +func (c *cancelledToolContext) Deadline() (deadline time.Time, ok bool) { + return c.cancelCtx.Deadline() +} + +func (c *cancelledToolContext) Value(key any) any { + return c.cancelCtx.Value(key) +} + // handleFunctionCalls calls the functions and returns the function response event. // // TODO: accept filters to include/exclude function calls. // TODO: check feasibility of running tool.Run concurrently. -func (f *Flow) handleFunctionCalls(ctx agent.InvocationContext, toolsDict map[string]tool.Tool, resp *model.LLMResponse, toolConfirmations map[string]*toolconfirmation.ToolConfirmation) (mergedEvent *session.Event, err error) { +func (f *Flow) handleFunctionCalls(ctx agent.InvocationContext, toolsDict map[string]tool.Tool, resp *model.LLMResponse, toolConfirmations map[string]*toolconfirmation.ToolConfirmation, liveSess agent.LiveSession) (mergedEvent *session.Event, err error) { fnCalls := utils.FunctionCalls(resp.Content) toolNames := slices.Collect(maps.Keys(toolsDict)) @@ -619,28 +1059,96 @@ func (f *Flow) handleFunctionCalls(ctx agent.InvocationContext, toolsDict map[st if toolConfirmations != nil { confirmation = toolConfirmations[fnCall.ID] } - toolCtx := toolinternal.NewToolContext(toolCallCtx, fnCall.ID, &session.EventActions{StateDelta: make(map[string]any)}, confirmation) + toolCtx := agent.NewToolContext(toolCallCtx, fnCall.ID, &session.EventActions{StateDelta: make(map[string]any)}, confirmation) var result map[string]any - curTool, found := toolsDict[fnCall.Name] - if !found { - err := newToolNotFoundError(fnCall.Name, toolNames) - result, err = f.runOnToolErrorCallbacks(toolCtx, &fakeTool{name: fnCall.Name}, fnCall.Args, err) - if err != nil { - result = map[string]any{"error": err.Error()} - } - } else if funcTool, ok := curTool.(toolinternal.FunctionTool); !ok { - err := newToolNotFoundError(fnCall.Name, toolNames) - result, err = f.runOnToolErrorCallbacks(toolCtx, &fakeTool{name: fnCall.Name}, fnCall.Args, err) - if err != nil { - result = map[string]any{"error": err.Error()} + var curTool tool.Tool + if fnCall.Name == "stop_streaming" { + funcToStop, _ := fnCall.Args["function_name"].(string) + var status string + if impl, ok := liveSess.(*liveSessionImpl); ok && impl.CancelAllStreamingTools(funcToStop) { + status = fmt.Sprintf("Successfully stopped all running instances of %s", funcToStop) + } else { + status = fmt.Sprintf("No active streaming function named %s found", funcToStop) } + result = map[string]any{"status": status} } else { - result = f.callTool(toolCtx, funcTool, fnCall.Args) + var found bool + curTool, found = toolsDict[fnCall.Name] + if !found { + err := newToolNotFoundError(fnCall.Name, toolNames) + result, err = f.runOnToolErrorCallbacks(toolCtx, &fakeTool{name: fnCall.Name}, fnCall.Args, err) + if err != nil { + result = map[string]any{"error": err.Error()} + } + } else if streamTool, ok := curTool.(toolinternal.StreamingFunctionTool); ok { + if liveSess != nil { + result = map[string]any{"status": "The function is running asynchronously and the results are pending."} + cancelCtx, cancel := context.WithCancel(toolCtx) + cancelToolCtx := &cancelledToolContext{ + ToolContext: toolCtx, + cancelCtx: cancelCtx, + } + if impl, ok := liveSess.(*liveSessionImpl); ok { + impl.RegisterStreamingTool(streamTool.Name(), fnCall.ID, cancel) + } + go func() { + defer func() { + if impl, ok := liveSess.(*liveSessionImpl); ok { + impl.UnregisterStreamingTool(streamTool.Name(), fnCall.ID) + } + cancel() + }() + for chunk, err := range streamTool.RunStream(cancelToolCtx, fnCall.Args) { + select { + case <-cancelCtx.Done(): + return + default: + } + if err != nil { + fmt.Printf("Error in streaming tool %s: %v\n", streamTool.Name(), err) + return + } + updatedContent := &genai.Content{ + Role: "user", + Parts: []*genai.Part{ + { + Text: fmt.Sprintf("Function %s returned: %s", streamTool.Name(), chunk), + }, + }, + } + if err := liveSess.Send(agent.LiveRequest{Content: updatedContent}); err != nil { + fmt.Printf("Failed to send content from streaming tool %s: %v\n", streamTool.Name(), err) + return + } + } + }() + } else { + var sb strings.Builder + for chunk, err := range streamTool.RunStream(toolCtx, fnCall.Args) { + if err != nil { + result = map[string]any{"error": err.Error()} + break + } + sb.WriteString(chunk) + } + if result == nil { + result = map[string]any{"result": sb.String()} + } + } + } else if funcTool, ok := curTool.(toolinternal.FunctionTool); !ok { + err := newToolNotFoundError(fnCall.Name, toolNames) + result, err = f.runOnToolErrorCallbacks(toolCtx, &fakeTool{name: fnCall.Name}, fnCall.Args, err) + if err != nil { + result = map[string]any{"error": err.Error()} + } + } else { + result = f.callTool(toolCtx, funcTool, fnCall.Args) + } } // TODO: handle long-running tool. - ev := session.NewEvent(ctx.InvocationID()) + ev := session.NewEventWithContext(ctx, ctx.InvocationID()) ev.LLMResponse = model.LLMResponse{ Content: &genai.Content{ Role: "user", @@ -689,7 +1197,7 @@ func (f *Flow) handleFunctionCalls(ctx agent.InvocationContext, toolsDict map[st return mergedEvent, nil } -func (f *Flow) runOnToolErrorCallbacks(toolCtx tool.Context, tool tool.Tool, fArgs map[string]any, err error) (map[string]any, error) { +func (f *Flow) runOnToolErrorCallbacks(toolCtx agent.ToolContext, tool tool.Tool, fArgs map[string]any, err error) (map[string]any, error) { pluginManager := pluginManagerFromContext(toolCtx) if pluginManager != nil { result, err := pluginManager.RunOnToolErrorCallback(toolCtx, tool, fArgs, err) @@ -700,7 +1208,7 @@ func (f *Flow) runOnToolErrorCallbacks(toolCtx tool.Context, tool tool.Tool, fAr return f.invokeOnToolErrorCallbacks(toolCtx, tool, fArgs, err) } -func (f *Flow) callTool(toolCtx tool.Context, tool toolinternal.FunctionTool, fArgs map[string]any) map[string]any { +func (f *Flow) callTool(toolCtx agent.ToolContext, tool toolinternal.FunctionTool, fArgs map[string]any) map[string]any { var response map[string]any var err error pluginManager := pluginManagerFromContext(toolCtx) @@ -747,7 +1255,7 @@ func (f *Flow) callTool(toolCtx tool.Context, tool toolinternal.FunctionTool, fA return response } -func (f *Flow) invokeBeforeToolCallbacks(toolCtx tool.Context, tool tool.Tool, fArgs map[string]any) (map[string]any, error) { +func (f *Flow) invokeBeforeToolCallbacks(toolCtx agent.ToolContext, tool tool.Tool, fArgs map[string]any) (map[string]any, error) { for _, callback := range f.BeforeToolCallbacks { result, err := callback(toolCtx, tool, fArgs) if err != nil { @@ -762,7 +1270,7 @@ func (f *Flow) invokeBeforeToolCallbacks(toolCtx tool.Context, tool tool.Tool, f return nil, nil } -func (f *Flow) invokeAfterToolCallbacks(toolCtx tool.Context, tool toolinternal.FunctionTool, fArgs, fResult map[string]any, fErr error) (map[string]any, error) { +func (f *Flow) invokeAfterToolCallbacks(toolCtx agent.ToolContext, tool toolinternal.FunctionTool, fArgs, fResult map[string]any, fErr error) (map[string]any, error) { for _, callback := range f.AfterToolCallbacks { result, err := callback(toolCtx, tool, fArgs, fResult, fErr) if err != nil { @@ -778,7 +1286,7 @@ func (f *Flow) invokeAfterToolCallbacks(toolCtx tool.Context, tool toolinternal. return fResult, fErr } -func (f *Flow) invokeOnToolErrorCallbacks(toolCtx tool.Context, tool tool.Tool, fArgs map[string]any, fErr error) (map[string]any, error) { +func (f *Flow) invokeOnToolErrorCallbacks(toolCtx agent.ToolContext, tool tool.Tool, fArgs map[string]any, fErr error) (map[string]any, error) { for _, callback := range f.OnToolErrorCallbacks { result, err := callback(toolCtx, tool, fArgs, fErr) if err != nil { @@ -880,7 +1388,7 @@ type pluginManager interface { RunBeforeModelCallback(cctx agent.CallbackContext, llmRequest *model.LLMRequest) (*model.LLMResponse, error) RunAfterModelCallback(cctx agent.CallbackContext, llmResponse *model.LLMResponse, llmResponseError error) (*model.LLMResponse, error) RunOnModelErrorCallback(ctx agent.CallbackContext, llmRequest *model.LLMRequest, llmResponseError error) (*model.LLMResponse, error) - RunBeforeToolCallback(ctx tool.Context, t tool.Tool, args map[string]any) (map[string]any, error) - RunAfterToolCallback(ctx tool.Context, t tool.Tool, args, result map[string]any, err error) (map[string]any, error) - RunOnToolErrorCallback(ctx tool.Context, t tool.Tool, args map[string]any, err error) (map[string]any, error) + RunBeforeToolCallback(ctx agent.ToolContext, t tool.Tool, args map[string]any) (map[string]any, error) + RunAfterToolCallback(ctx agent.ToolContext, t tool.Tool, args, result map[string]any, err error) (map[string]any, error) + RunOnToolErrorCallback(ctx agent.ToolContext, t tool.Tool, args map[string]any, err error) (map[string]any, error) } diff --git a/internal/llminternal/base_flow_test.go b/internal/llminternal/base_flow_test.go index 3b211df9c..6d18e43d8 100644 --- a/internal/llminternal/base_flow_test.go +++ b/internal/llminternal/base_flow_test.go @@ -31,7 +31,7 @@ import ( type mockFunctionTool struct { name string - runFunc func(tool.Context, map[string]any) (map[string]any, error) + runFunc func(agent.ToolContext, map[string]any) (map[string]any, error) } func (m *mockFunctionTool) Name() string { @@ -54,11 +54,11 @@ func (m *mockFunctionTool) IsLongRunning() bool { return false } -func (m *mockFunctionTool) ProcessRequest(ctx tool.Context, req *model.LLMRequest) error { +func (m *mockFunctionTool) ProcessRequest(ctx agent.ToolContext, req *model.LLMRequest) error { return nil } -func (m *mockFunctionTool) Run(ctx tool.Context, args any) (map[string]any, error) { +func (m *mockFunctionTool) Run(ctx agent.ToolContext, args any) (map[string]any, error) { if m.runFunc != nil { return m.runFunc(ctx, args.(map[string]any)) } @@ -80,10 +80,10 @@ func (m *mockToolset) Tools(ctx agent.ReadonlyContext) ([]tool.Tool, error) { type mockRequestProcessorToolset struct { name string - process func(ctx tool.Context, req *model.LLMRequest) error + process func(ctx agent.ToolContext, req *model.LLMRequest) error } -func (m *mockRequestProcessorToolset) ProcessRequest(ctx tool.Context, req *model.LLMRequest) error { +func (m *mockRequestProcessorToolset) ProcessRequest(ctx agent.ToolContext, req *model.LLMRequest) error { if m.process != nil { return m.process(ctx, req) } @@ -110,7 +110,7 @@ func TestCallTool(t *testing.T) { name: "tool runs successfully", tool: &mockFunctionTool{ name: "testTool", - runFunc: func(ctx tool.Context, args map[string]any) (map[string]any, error) { + runFunc: func(ctx agent.ToolContext, args map[string]any) (map[string]any, error) { return map[string]any{"result": "success"}, nil }, }, @@ -121,7 +121,7 @@ func TestCallTool(t *testing.T) { name: "tool error", tool: &mockFunctionTool{ name: "testTool", - runFunc: func(ctx tool.Context, args map[string]any) (map[string]any, error) { + runFunc: func(ctx agent.ToolContext, args map[string]any) (map[string]any, error) { return nil, errors.New("tool error") }, }, @@ -132,16 +132,16 @@ func TestCallTool(t *testing.T) { name: "before callback returns result", tool: &mockFunctionTool{ name: "testTool", - runFunc: func(ctx tool.Context, args map[string]any) (map[string]any, error) { + runFunc: func(ctx agent.ToolContext, args map[string]any) (map[string]any, error) { t.Error("tool should not be called") return nil, nil }, }, beforeToolCallbacks: []BeforeToolCallback{ - func(ctx tool.Context, tool tool.Tool, args map[string]any) (map[string]any, error) { + func(ctx agent.ToolContext, tool tool.Tool, args map[string]any) (map[string]any, error) { return map[string]any{"result": "intercepted"}, nil }, - func(ctx tool.Context, tool tool.Tool, args map[string]any) (map[string]any, error) { + func(ctx agent.ToolContext, tool tool.Tool, args map[string]any) (map[string]any, error) { return map[string]any{"result": "2nd callback should not be called"}, nil }, }, @@ -151,16 +151,16 @@ func TestCallTool(t *testing.T) { name: "before callback returns error", tool: &mockFunctionTool{ name: "testTool", - runFunc: func(ctx tool.Context, args map[string]any) (map[string]any, error) { + runFunc: func(ctx agent.ToolContext, args map[string]any) (map[string]any, error) { t.Error("tool should not be called") return nil, nil }, }, beforeToolCallbacks: []BeforeToolCallback{ - func(ctx tool.Context, tool tool.Tool, args map[string]any) (map[string]any, error) { + func(ctx agent.ToolContext, tool tool.Tool, args map[string]any) (map[string]any, error) { return nil, errors.New("before callback error") }, - func(ctx tool.Context, tool tool.Tool, args map[string]any) (map[string]any, error) { + func(ctx agent.ToolContext, tool tool.Tool, args map[string]any) (map[string]any, error) { return nil, errors.New("unexpected error") }, }, @@ -170,12 +170,12 @@ func TestCallTool(t *testing.T) { name: "after callback modifies result", tool: &mockFunctionTool{ name: "testTool", - runFunc: func(ctx tool.Context, args map[string]any) (map[string]any, error) { + runFunc: func(ctx agent.ToolContext, args map[string]any) (map[string]any, error) { return map[string]any{"result": "original"}, nil }, }, afterToolCallbacks: []AfterToolCallback{ - func(ctx tool.Context, tool tool.Tool, args, result map[string]any, err error) (map[string]any, error) { + func(ctx agent.ToolContext, tool tool.Tool, args, result map[string]any, err error) (map[string]any, error) { return map[string]any{"result": "modified"}, nil }, }, @@ -185,18 +185,18 @@ func TestCallTool(t *testing.T) { name: "after callback handles error", tool: &mockFunctionTool{ name: "testTool", - runFunc: func(ctx tool.Context, args map[string]any) (map[string]any, error) { + runFunc: func(ctx agent.ToolContext, args map[string]any) (map[string]any, error) { return nil, errors.New("tool error") }, }, afterToolCallbacks: []AfterToolCallback{ - func(ctx tool.Context, tool tool.Tool, args, result map[string]any, err error) (map[string]any, error) { + func(ctx agent.ToolContext, tool tool.Tool, args, result map[string]any, err error) (map[string]any, error) { if err != nil { return map[string]any{"result": "error handled"}, nil } return nil, nil }, - func(ctx tool.Context, tool tool.Tool, args, result map[string]any, err error) (map[string]any, error) { + func(ctx agent.ToolContext, tool tool.Tool, args, result map[string]any, err error) (map[string]any, error) { return map[string]any{"result": "unexpected output"}, nil }, }, @@ -206,15 +206,15 @@ func TestCallTool(t *testing.T) { name: "after callback returns error", tool: &mockFunctionTool{ name: "testTool", - runFunc: func(ctx tool.Context, args map[string]any) (map[string]any, error) { + runFunc: func(ctx agent.ToolContext, args map[string]any) (map[string]any, error) { return map[string]any{"result": "success"}, nil }, }, afterToolCallbacks: []AfterToolCallback{ - func(ctx tool.Context, tool tool.Tool, args, result map[string]any, err error) (map[string]any, error) { + func(ctx agent.ToolContext, tool tool.Tool, args, result map[string]any, err error) (map[string]any, error) { return nil, errors.New("after callback error") }, - func(ctx tool.Context, tool tool.Tool, args, result map[string]any, err error) (map[string]any, error) { + func(ctx agent.ToolContext, tool tool.Tool, args, result map[string]any, err error) (map[string]any, error) { return nil, errors.New("unexpected error") }, }, @@ -224,17 +224,17 @@ func TestCallTool(t *testing.T) { name: "no-op callbacks return func results", tool: &mockFunctionTool{ name: "testTool", - runFunc: func(ctx tool.Context, args map[string]any) (map[string]any, error) { + runFunc: func(ctx agent.ToolContext, args map[string]any) (map[string]any, error) { return map[string]any{"result": "success"}, nil }, }, beforeToolCallbacks: []BeforeToolCallback{ - func(ctx tool.Context, tool tool.Tool, args map[string]any) (map[string]any, error) { + func(ctx agent.ToolContext, tool tool.Tool, args map[string]any) (map[string]any, error) { return nil, nil }, }, afterToolCallbacks: []AfterToolCallback{ - func(ctx tool.Context, tool tool.Tool, args, result map[string]any, err error) (map[string]any, error) { + func(ctx agent.ToolContext, tool tool.Tool, args, result map[string]any, err error) (map[string]any, error) { return nil, nil }, }, @@ -244,18 +244,18 @@ func TestCallTool(t *testing.T) { name: "before callback result passed to after callback", tool: &mockFunctionTool{ name: "testTool", - runFunc: func(ctx tool.Context, args map[string]any) (map[string]any, error) { + runFunc: func(ctx agent.ToolContext, args map[string]any) (map[string]any, error) { t.Error("tool should not be called") return nil, nil }, }, beforeToolCallbacks: []BeforeToolCallback{ - func(ctx tool.Context, tool tool.Tool, args map[string]any) (map[string]any, error) { + func(ctx agent.ToolContext, tool tool.Tool, args map[string]any) (map[string]any, error) { return map[string]any{"result": "from_before"}, nil }, }, afterToolCallbacks: []AfterToolCallback{ - func(ctx tool.Context, tool tool.Tool, args, result map[string]any, err error) (map[string]any, error) { + func(ctx agent.ToolContext, tool tool.Tool, args, result map[string]any, err error) (map[string]any, error) { if val, ok := result["result"]; !ok || val != "from_before" { return nil, errors.New("unexpected result in after callback") } @@ -268,18 +268,18 @@ func TestCallTool(t *testing.T) { name: "before callback error passed to after callback", tool: &mockFunctionTool{ name: "testTool", - runFunc: func(ctx tool.Context, args map[string]any) (map[string]any, error) { + runFunc: func(ctx agent.ToolContext, args map[string]any) (map[string]any, error) { t.Error("tool should not be called") return nil, nil }, }, beforeToolCallbacks: []BeforeToolCallback{ - func(ctx tool.Context, tool tool.Tool, args map[string]any) (map[string]any, error) { + func(ctx agent.ToolContext, tool tool.Tool, args map[string]any) (map[string]any, error) { return nil, errors.New("error_from_before") }, }, afterToolCallbacks: []AfterToolCallback{ - func(ctx tool.Context, tool tool.Tool, args, result map[string]any, err error) (map[string]any, error) { + func(ctx agent.ToolContext, tool tool.Tool, args, result map[string]any, err error) (map[string]any, error) { if err == nil || err.Error() != "error_from_before" { return nil, errors.New("unexpected error in after callback") } @@ -292,18 +292,18 @@ func TestCallTool(t *testing.T) { name: "before callback error passed to on tool error callback", tool: &mockFunctionTool{ name: "testTool", - runFunc: func(ctx tool.Context, args map[string]any) (map[string]any, error) { + runFunc: func(ctx agent.ToolContext, args map[string]any) (map[string]any, error) { t.Error("tool should not be called") return nil, nil }, }, beforeToolCallbacks: []BeforeToolCallback{ - func(ctx tool.Context, tool tool.Tool, args map[string]any) (map[string]any, error) { + func(ctx agent.ToolContext, tool tool.Tool, args map[string]any) (map[string]any, error) { return nil, errors.New("error_from_before") }, }, onToolErrorCallbacks: []OnToolErrorCallback{ - func(ctx tool.Context, tool tool.Tool, args map[string]any, err error) (map[string]any, error) { + func(ctx agent.ToolContext, tool tool.Tool, args map[string]any, err error) (map[string]any, error) { if err == nil || err.Error() != "error_from_before" { t.Error("unexpected error in on tool error callback") return nil, errors.New("unexpected error in on tool error callback") @@ -317,18 +317,18 @@ func TestCallTool(t *testing.T) { name: "before callback error passed to on tool error callback and after tool called", tool: &mockFunctionTool{ name: "testTool", - runFunc: func(ctx tool.Context, args map[string]any) (map[string]any, error) { + runFunc: func(ctx agent.ToolContext, args map[string]any) (map[string]any, error) { t.Error("tool should not be called") return nil, nil }, }, beforeToolCallbacks: []BeforeToolCallback{ - func(ctx tool.Context, tool tool.Tool, args map[string]any) (map[string]any, error) { + func(ctx agent.ToolContext, tool tool.Tool, args map[string]any) (map[string]any, error) { return nil, errors.New("error_from_before") }, }, onToolErrorCallbacks: []OnToolErrorCallback{ - func(ctx tool.Context, tool tool.Tool, args map[string]any, err error) (map[string]any, error) { + func(ctx agent.ToolContext, tool tool.Tool, args map[string]any, err error) (map[string]any, error) { if err == nil || err.Error() != "error_from_before" { t.Error("unexpected error in on tool error callback") return nil, errors.New("unexpected error in on tool error callback") @@ -337,7 +337,7 @@ func TestCallTool(t *testing.T) { }, }, afterToolCallbacks: []AfterToolCallback{ - func(ctx tool.Context, tool tool.Tool, args, result map[string]any, err error) (map[string]any, error) { + func(ctx agent.ToolContext, tool tool.Tool, args, result map[string]any, err error) (map[string]any, error) { if err != nil { return nil, errors.New("unexpected error in after callback") } @@ -350,18 +350,18 @@ func TestCallTool(t *testing.T) { name: "before callback error passed to on tool error callback and passed to after tool called", tool: &mockFunctionTool{ name: "testTool", - runFunc: func(ctx tool.Context, args map[string]any) (map[string]any, error) { + runFunc: func(ctx agent.ToolContext, args map[string]any) (map[string]any, error) { t.Error("tool should not be called") return nil, nil }, }, beforeToolCallbacks: []BeforeToolCallback{ - func(ctx tool.Context, tool tool.Tool, args map[string]any) (map[string]any, error) { + func(ctx agent.ToolContext, tool tool.Tool, args map[string]any) (map[string]any, error) { return nil, errors.New("error_from_before") }, }, onToolErrorCallbacks: []OnToolErrorCallback{ - func(ctx tool.Context, tool tool.Tool, args map[string]any, err error) (map[string]any, error) { + func(ctx agent.ToolContext, tool tool.Tool, args map[string]any, err error) (map[string]any, error) { if err == nil || err.Error() != "error_from_before" { t.Error("unexpected error in on tool error callback") return nil, errors.New("unexpected error in on tool error callback") @@ -370,7 +370,7 @@ func TestCallTool(t *testing.T) { }, }, afterToolCallbacks: []AfterToolCallback{ - func(ctx tool.Context, tool tool.Tool, args, result map[string]any, err error) (map[string]any, error) { + func(ctx agent.ToolContext, tool tool.Tool, args, result map[string]any, err error) (map[string]any, error) { if err == nil || err.Error() != "error_from_on_tool_error" { return nil, errors.New("unexpected error in after callback") } @@ -383,18 +383,18 @@ func TestCallTool(t *testing.T) { name: "before callback error passed to on tool error callback and passed to after tool called and handled", tool: &mockFunctionTool{ name: "testTool", - runFunc: func(ctx tool.Context, args map[string]any) (map[string]any, error) { + runFunc: func(ctx agent.ToolContext, args map[string]any) (map[string]any, error) { t.Error("tool should not be called") return nil, nil }, }, beforeToolCallbacks: []BeforeToolCallback{ - func(ctx tool.Context, tool tool.Tool, args map[string]any) (map[string]any, error) { + func(ctx agent.ToolContext, tool tool.Tool, args map[string]any) (map[string]any, error) { return nil, errors.New("error_from_before") }, }, onToolErrorCallbacks: []OnToolErrorCallback{ - func(ctx tool.Context, tool tool.Tool, args map[string]any, err error) (map[string]any, error) { + func(ctx agent.ToolContext, tool tool.Tool, args map[string]any, err error) (map[string]any, error) { if err == nil || err.Error() != "error_from_before" { t.Error("unexpected error in on tool error callback") return nil, errors.New("unexpected error in on tool error callback") @@ -403,7 +403,7 @@ func TestCallTool(t *testing.T) { }, }, afterToolCallbacks: []AfterToolCallback{ - func(ctx tool.Context, tool tool.Tool, args, result map[string]any, err error) (map[string]any, error) { + func(ctx agent.ToolContext, tool tool.Tool, args, result map[string]any, err error) (map[string]any, error) { if err == nil || err.Error() != "error_from_on_tool_error" { return nil, errors.New("unexpected error in after callback") } @@ -421,7 +421,7 @@ func TestCallTool(t *testing.T) { OnToolErrorCallbacks: tc.onToolErrorCallbacks, } ctx := icontext.NewInvocationContext(t.Context(), icontext.InvocationContextParams{}) - got := f.callTool(toolinternal.NewToolContext(ctx, "", nil, nil), tc.tool, tc.args) + got := f.callTool(agent.NewToolContext(ctx, "", nil, nil), tc.tool, tc.args) if diff := cmp.Diff(tc.want, got); diff != "" { t.Errorf("callTool() mismatch (-want +got):\n%s", diff) } @@ -630,7 +630,7 @@ func TestPreprocess_Toolset(t *testing.T) { s: &State{ Toolsets: []tool.Toolset{&mockRequestProcessorToolset{ name: "toolset", - process: func(_ tool.Context, _ *model.LLMRequest) error { + process: func(_ agent.ToolContext, _ *model.LLMRequest) error { return errors.New("process error") }, }}, @@ -646,7 +646,7 @@ func TestPreprocess_Toolset(t *testing.T) { &mockToolset{name: "toolset_without_processor"}, &mockRequestProcessorToolset{ name: "toolset_with_processor", - process: func(_ tool.Context, req *model.LLMRequest) error { + process: func(_ agent.ToolContext, req *model.LLMRequest) error { req.Model = "modified-model" return nil }, diff --git a/internal/llminternal/contents_processor.go b/internal/llminternal/contents_processor.go index 360f953e7..46fb5a141 100644 --- a/internal/llminternal/contents_processor.go +++ b/internal/llminternal/contents_processor.go @@ -70,9 +70,9 @@ func buildContentsDefault(agentName, invocationBranch string, events []*session. for _, ev := range events { content := utils.Content(ev) // Skip events without content or generated neither by user nor - // by model. - // e.g. events purely for mutating session states. - if content == nil || content.Role == "" || len(content.Parts) == 0 { + // by model, UNLESS they have transcriptions. + if (content == nil || content.Role == "" || len(content.Parts) == 0) && + ev.LLMResponse.InputTranscription == nil && ev.LLMResponse.OutputTranscription == nil { // TODO: log a bad event with content but no Role is skipped // Note: python checks here if content.Parts[0] is an empty string and skip if so. // But unlike python that distinguishes None vs empty string, two cases are indistinguishable in Go. @@ -93,6 +93,53 @@ func buildContentsDefault(agentName, invocationBranch string, events []*session. } } + // Aggregate transcription events (convert to text parts on the fly) + var processedEvents []*session.Event + var accumulatedInputTranscription string + var accumulatedOutputTranscription string + + for i := 0; i < len(filtered); i++ { + ev := filtered[i] + content := utils.Content(ev) + if content == nil || len(content.Parts) == 0 { + if ev.LLMResponse.InputTranscription != nil && ev.LLMResponse.InputTranscription.Text != "" { + accumulatedInputTranscription += ev.LLMResponse.InputTranscription.Text + if i != len(filtered)-1 && + filtered[i+1].LLMResponse.InputTranscription != nil && + filtered[i+1].LLMResponse.InputTranscription.Text != "" { + continue + } + // Create a new event with content + newEv := cloneEvent(ev) + newEv.LLMResponse.InputTranscription = nil + newEv.LLMResponse.Content = &genai.Content{ + Role: genai.RoleUser, + Parts: []*genai.Part{{Text: accumulatedInputTranscription}}, + } + ev = newEv + accumulatedInputTranscription = "" + } else if ev.LLMResponse.OutputTranscription != nil && ev.LLMResponse.OutputTranscription.Text != "" { + accumulatedOutputTranscription += ev.LLMResponse.OutputTranscription.Text + if i != len(filtered)-1 && + filtered[i+1].LLMResponse.OutputTranscription != nil && + filtered[i+1].LLMResponse.OutputTranscription.Text != "" { + continue + } + // Create a new event with content + newEv := cloneEvent(ev) + newEv.LLMResponse.OutputTranscription = nil + newEv.LLMResponse.Content = &genai.Content{ + Role: "model", + Parts: []*genai.Part{{Text: accumulatedOutputTranscription}}, + } + ev = newEv + accumulatedOutputTranscription = "" + } + } + processedEvents = append(processedEvents, ev) + } + filtered = processedEvents + // src/google/adk/flows/llm_flows/contents.py // - _rearrange_events_for_async_function_response filtered, err := rearrangeEventsForLatestFunctionResponse(filtered) @@ -225,10 +272,18 @@ SearchLoop: // A label to allow breaking out of the nested loop ) } - // Collect all function response events *between* the call and the last response. + // Collect function response events related to the matching call while + // preserving unrelated tool events that happened in between. var responseEventsToMerge []*session.Event + resultEvents := events[:functionCallEventIdx+1] for i := functionCallEventIdx + 1; i < len(events)-1; i++ { event := events[i] + calls := utils.FunctionCalls(event.Content) + if len(calls) > 0 { + resultEvents = append(resultEvents, event) + continue + } + responses := utils.FunctionResponses(event.Content) if len(responses) == 0 { continue @@ -245,13 +300,14 @@ SearchLoop: // A label to allow breaking out of the nested loop if isRelated { responseEventsToMerge = append(responseEventsToMerge, event) + } else { + resultEvents = append(resultEvents, event) } } // Add the final response event itself to the list to be merged. responseEventsToMerge = append(responseEventsToMerge, events[len(events)-1]) - resultEvents := events[:functionCallEventIdx+1] mergedEvent, err := mergeFunctionResponseEvents(responseEventsToMerge) if err != nil { return nil, err diff --git a/internal/llminternal/contents_processor_test.go b/internal/llminternal/contents_processor_test.go index ac5dbbdb3..3fce0d8d7 100644 --- a/internal/llminternal/contents_processor_test.go +++ b/internal/llminternal/contents_processor_test.go @@ -31,6 +31,7 @@ import ( "google.golang.org/adk/internal/utils" "google.golang.org/adk/model" "google.golang.org/adk/session" + "google.golang.org/adk/tool/toolconfirmation" ) type testModel struct { @@ -418,6 +419,46 @@ func TestContentsRequestProcessor(t *testing.T) { }, want: nil, }, + { + name: "TranscriptionAggregation", + events: []*session.Event{ + { + Author: "user", + LLMResponse: model.LLMResponse{ + InputTranscription: &genai.Transcription{Text: "hello ", Finished: false}, + }, + }, + { + Author: "user", + LLMResponse: model.LLMResponse{ + InputTranscription: &genai.Transcription{Text: "world", Finished: true}, + }, + }, + { + Author: "testAgent", + LLMResponse: model.LLMResponse{ + OutputTranscription: &genai.Transcription{Text: "hi ", Finished: false}, + }, + }, + { + Author: "testAgent", + LLMResponse: model.LLMResponse{ + OutputTranscription: &genai.Transcription{Text: "there", Finished: true}, + }, + }, + { + Author: "user", + LLMResponse: model.LLMResponse{ + InputTranscription: &genai.Transcription{Text: "ok", Finished: true}, + }, + }, + }, + want: []*genai.Content{ + genai.NewContentFromText("hello world", "user"), + genai.NewContentFromText("hi there", "model"), + genai.NewContentFromText("ok", "user"), + }, + }, } for _, tc := range testCases { @@ -705,6 +746,70 @@ func TestContentsRequestProcessor_Rearrange(t *testing.T) { Response: map[string]any{"output": "preserved"}, } + // Human-in-the-loop confirmation call/responses + fcHITLApproved := &genai.FunctionCall{ + ID: "hitl_approved_call", + Name: "request_vacation_days", + Args: map[string]any{"days": 5, "user_id": "user-123"}, + } + frHITLApprovedPending := &genai.FunctionResponse{ + ID: "hitl_approved_call", + Name: "request_vacation_days", + Response: map[string]any{"status": "Manager approval is required.", "request_id": "req-1"}, + } + frHITLApprovedFinal := &genai.FunctionResponse{ + ID: "hitl_approved_call", + Name: "request_vacation_days", + Response: map[string]any{"status": "The time off request is accepted.", "days_approved": 5, "request_id": "req-1"}, + } + fcHITLApprovalRequest := &genai.FunctionCall{ + ID: "hitl_approved_confirmation", + Name: toolconfirmation.FunctionCallName, + Args: map[string]any{ + "originalFunctionCall": fcHITLApproved, + "toolConfirmation": toolconfirmation.ToolConfirmation{ + Confirmed: false, + Hint: "Please approve or reject the tool call request_vacation_days().", + }, + }, + } + frHITLApprovalConfirmed := &genai.FunctionResponse{ + ID: "hitl_approved_confirmation", + Name: toolconfirmation.FunctionCallName, + Response: map[string]any{"confirmed": true}, + } + fcHITLDenied := &genai.FunctionCall{ + ID: "hitl_denied_call", + Name: "request_vacation_days", + Args: map[string]any{"days": 10, "user_id": "user-123"}, + } + frHITLDeniedPending := &genai.FunctionResponse{ + ID: "hitl_denied_call", + Name: "request_vacation_days", + Response: map[string]any{"status": "Manager approval is required.", "request_id": "req-2"}, + } + frHITLDeniedFinal := &genai.FunctionResponse{ + ID: "hitl_denied_call", + Name: "request_vacation_days", + Response: map[string]any{"status": "The time off request is rejected.", "days_approved": 0, "request_id": "req-2"}, + } + fcHITLDenialRequest := &genai.FunctionCall{ + ID: "hitl_denied_confirmation", + Name: toolconfirmation.FunctionCallName, + Args: map[string]any{ + "originalFunctionCall": fcHITLDenied, + "toolConfirmation": toolconfirmation.ToolConfirmation{ + Confirmed: false, + Hint: "Please approve or reject the tool call request_vacation_days().", + }, + }, + } + frHITLDenialRejected := &genai.FunctionResponse{ + ID: "hitl_denied_confirmation", + Name: toolconfirmation.FunctionCallName, + Response: map[string]any{"confirmed": false}, + } + // Error Call/Response frOrphaned := &genai.FunctionResponse{ ID: "no_matching_call", @@ -764,6 +869,56 @@ func TestContentsRequestProcessor_Rearrange(t *testing.T) { NewContentFromFunctionResponse(frLROFinal, "user"), }, }, + { + name: "Rearrangement preserves unrelated function events", + events: []*session.Event{ + {Author: "user", LLMResponse: model.LLMResponse{Content: genai.NewContentFromText("Run long process and search", "user")}}, + {Author: agentName, LLMResponse: model.LLMResponse{Content: NewContentFromFunctionCall(fcLRO, "model")}}, + {Author: "user", LLMResponse: model.LLMResponse{Content: NewContentFromFunctionResponse(frLROInter, "user")}}, + {Author: agentName, LLMResponse: model.LLMResponse{Content: NewContentFromFunctionCall(fcBasic, "model")}}, + {Author: "user", LLMResponse: model.LLMResponse{Content: NewContentFromFunctionResponse(frBasic, "user")}}, + {Author: "user", LLMResponse: model.LLMResponse{Content: NewContentFromFunctionResponse(frLROFinal, "user")}}, + }, + want: []*genai.Content{ + genai.NewContentFromText("Run long process and search", "user"), + NewContentFromFunctionCall(fcLRO, "model"), + NewContentFromFunctionResponse(frLROFinal, "user"), + NewContentFromFunctionCall(fcBasic, "model"), + NewContentFromFunctionResponse(frBasic, "user"), + }, + }, + { + name: "HITL confirmation approved preserves resumed tool response", + events: []*session.Event{ + {Author: "user", LLMResponse: model.LLMResponse{Content: genai.NewContentFromText("Request five vacation days", "user")}}, + {Author: agentName, LLMResponse: model.LLMResponse{Content: NewContentFromFunctionCall(fcHITLApproved, "model")}}, + {Author: agentName, LLMResponse: model.LLMResponse{Content: NewContentFromFunctionResponse(frHITLApprovedPending, "user")}}, + {Author: agentName, LLMResponse: model.LLMResponse{Content: NewContentFromFunctionCall(fcHITLApprovalRequest, "model")}}, + {Author: "user", LLMResponse: model.LLMResponse{Content: NewContentFromFunctionResponse(frHITLApprovalConfirmed, "user")}}, + {Author: agentName, LLMResponse: model.LLMResponse{Content: NewContentFromFunctionResponse(frHITLApprovedFinal, "user")}}, + }, + want: []*genai.Content{ + genai.NewContentFromText("Request five vacation days", "user"), + NewContentFromFunctionCall(fcHITLApproved, "model"), + NewContentFromFunctionResponse(frHITLApprovedFinal, "user"), + }, + }, + { + name: "HITL confirmation denied preserves rejected tool response", + events: []*session.Event{ + {Author: "user", LLMResponse: model.LLMResponse{Content: genai.NewContentFromText("Request ten vacation days", "user")}}, + {Author: agentName, LLMResponse: model.LLMResponse{Content: NewContentFromFunctionCall(fcHITLDenied, "model")}}, + {Author: agentName, LLMResponse: model.LLMResponse{Content: NewContentFromFunctionResponse(frHITLDeniedPending, "user")}}, + {Author: agentName, LLMResponse: model.LLMResponse{Content: NewContentFromFunctionCall(fcHITLDenialRequest, "model")}}, + {Author: "user", LLMResponse: model.LLMResponse{Content: NewContentFromFunctionResponse(frHITLDenialRejected, "user")}}, + {Author: agentName, LLMResponse: model.LLMResponse{Content: NewContentFromFunctionResponse(frHITLDeniedFinal, "user")}}, + }, + want: []*genai.Content{ + genai.NewContentFromText("Request ten vacation days", "user"), + NewContentFromFunctionCall(fcHITLDenied, "model"), + NewContentFromFunctionResponse(frHITLDeniedFinal, "user"), + }, + }, { name: "Rearrangement with mixed LRO and normal calls", events: []*session.Event{ diff --git a/internal/llminternal/converters/converters.go b/internal/llminternal/converters/converters.go index e9759447f..7a30fd1ac 100644 --- a/internal/llminternal/converters/converters.go +++ b/internal/llminternal/converters/converters.go @@ -56,9 +56,17 @@ func Genai2LLMResponse(res *genai.GenerateContentResponse) *model.LLMResponse { ModelVersion: res.ModelVersion, } } + // no candidates, no prompt feedback. + // Sometimes gemini-3* invoked via aiplatform (VertexAI) sends empty entries at the beginning. + // sample stream of SSE chunks (first 3): + // data: {"candidates": [{"content": {"role": "model","parts": [{"text": ""}]}}],"usageMetadata": {"trafficType": "ON_DEMAND"},"modelVersion": "gemini-3.1-flash-lite","createTime": "2026-05-28T09:40:03.380865Z","responseId": "REDACTED"} + // + // data: {"usageMetadata": {"trafficType": "ON_DEMAND"},"modelVersion": "gemini-3.1-flash-lite","createTime": "2026-05-28T09:40:03.380865Z","responseId": "REDACTED"} + // + // data: {"usageMetadata": {"trafficType": "ON_DEMAND"},"modelVersion": "gemini-3.1-flash-lite","createTime": "2026-05-28T09:40:03.380865Z","responseId": "REDACTED"} + // we should treat them as valid, empty responses and let the downstream to process usageMetadata return &model.LLMResponse{ - ErrorCode: "UNKNOWN_ERROR", - ErrorMessage: "Unknown error.", + Content: &genai.Content{Parts: []*genai.Part{}, Role: "model"}, UsageMetadata: usageMetadata, ModelVersion: res.ModelVersion, } diff --git a/internal/llminternal/functions.go b/internal/llminternal/functions.go index 98db0dfb0..9112732e9 100644 --- a/internal/llminternal/functions.go +++ b/internal/llminternal/functions.go @@ -26,7 +26,7 @@ import ( // generateRequestConfirmationEvent creates a new Event containing // adk_request_confirmation function calls based on the requested confirmations. -// NOTE: The trigger for this in ADK Go is usually a tool.Context.RequestConfirmation call, +// NOTE: The trigger for this in ADK Go is usually a agent.ToolContext.RequestConfirmation call, // not parsing a function_response_event like in the Python example. // This function assumes you have a list of confirmations to process. func generateRequestConfirmationEvent( @@ -43,31 +43,34 @@ func generateRequestConfirmationEvent( parts := []*genai.Part{} longRunningToolIDs := []string{} - functionCalls := make(map[string]*genai.FunctionCall, len(functionCallEvent.Content.Parts)) - for _, call := range utils.FunctionCalls(functionCallEvent.Content) { - functionCalls[call.ID] = call + functionCallParts := make(map[string]*genai.Part, len(functionCallEvent.Content.Parts)) + for _, part := range functionCallEvent.Content.Parts { + if part.FunctionCall != nil { + functionCallParts[part.FunctionCall.ID] = part + } } for funcID, confirmation := range functionResponseEvent.Actions.RequestedToolConfirmations { - originalFunctionCall, ok := functionCalls[funcID] - if !ok || originalFunctionCall == nil { + originalPart, ok := functionCallParts[funcID] + if !ok || originalPart.FunctionCall == nil { continue } // Prepare arguments for the adk_request_confirmation call args := map[string]any{ - "originalFunctionCall": originalFunctionCall, + "originalFunctionCall": originalPart.FunctionCall, "toolConfirmation": confirmation, } requestConfirmationFC := &genai.FunctionCall{ - ID: utils.GenerateFunctionCallID(), + ID: utils.GenerateFunctionCallID(invocationContext), Name: toolconfirmation.FunctionCallName, Args: args, } parts = append(parts, &genai.Part{ - FunctionCall: requestConfirmationFC, + FunctionCall: requestConfirmationFC, + ThoughtSignature: originalPart.ThoughtSignature, }) longRunningToolIDs = append(longRunningToolIDs, requestConfirmationFC.ID) } @@ -76,7 +79,7 @@ func generateRequestConfirmationEvent( return nil } - ev := session.NewEvent(invocationContext.InvocationID()) + ev := session.NewEventWithContext(invocationContext, invocationContext.InvocationID()) ev.Author = invocationContext.Agent().Name() ev.Branch = invocationContext.Branch() ev.LLMResponse = model.LLMResponse{ diff --git a/internal/llminternal/functions_test.go b/internal/llminternal/functions_test.go index b7c935613..e0cca8179 100644 --- a/internal/llminternal/functions_test.go +++ b/internal/llminternal/functions_test.go @@ -16,6 +16,7 @@ package llminternal import ( "testing" + "time" "github.com/google/go-cmp/cmp" "github.com/google/go-cmp/cmp/cmpopts" @@ -55,6 +56,11 @@ func (m *mockInvocationContext) Branch() string { return m.branch } +func (m *mockInvocationContext) Deadline() (time.Time, bool) { return time.Time{}, false } +func (m *mockInvocationContext) Done() <-chan struct{} { return nil } +func (m *mockInvocationContext) Err() error { return nil } +func (m *mockInvocationContext) Value(any) any { return nil } + func TestGenerateRequestConfirmationEvent(t *testing.T) { confirmingFunctionCall := &genai.FunctionCall{ ID: "call_1", @@ -252,3 +258,79 @@ func TestGenerateRequestConfirmationEventHasID(t *testing.T) { t.Errorf("expected InvocationID=\"inv_1\", got %q", got.InvocationID) } } + +func TestGenerateRequestConfirmationEventPreservesThoughtSignature(t *testing.T) { + thoughtSignature := []byte("test-thought-signature") + ctx := &mockInvocationContext{ + invocationID: "inv_1", + agentName: "agent_1", + } + functionCallEvent := &session.Event{ + LLMResponse: model.LLMResponse{ + Content: &genai.Content{ + Parts: []*genai.Part{ + { + ThoughtSignature: thoughtSignature, + FunctionCall: &genai.FunctionCall{ + ID: "call_1", + Name: "test_tool", + Args: map[string]any{"arg": "val"}, + }, + }, + }, + }, + }, + } + functionResponseEvent := &session.Event{ + Actions: session.EventActions{ + RequestedToolConfirmations: map[string]toolconfirmation.ToolConfirmation{ + "call_1": {Hint: "Are you sure?"}, + }, + }, + } + + got := generateRequestConfirmationEvent(ctx, functionCallEvent, functionResponseEvent) + if got == nil || got.Content == nil || len(got.Content.Parts) != 1 { + t.Fatalf("expected one confirmation part, got %#v", got) + } + if diff := cmp.Diff(thoughtSignature, got.Content.Parts[0].ThoughtSignature); diff != "" { + t.Errorf("ThoughtSignature mismatch (-want +got):\n%s", diff) + } +} + +func TestGenerateRequestConfirmationEventNoThoughtSignature(t *testing.T) { + ctx := &mockInvocationContext{ + invocationID: "inv_1", + agentName: "agent_1", + } + functionCallEvent := &session.Event{ + LLMResponse: model.LLMResponse{ + Content: &genai.Content{ + Parts: []*genai.Part{ + { + FunctionCall: &genai.FunctionCall{ + ID: "call_1", + Name: "test_tool", + Args: map[string]any{"arg": "val"}, + }, + }, + }, + }, + }, + } + functionResponseEvent := &session.Event{ + Actions: session.EventActions{ + RequestedToolConfirmations: map[string]toolconfirmation.ToolConfirmation{ + "call_1": {Hint: "Are you sure?"}, + }, + }, + } + + got := generateRequestConfirmationEvent(ctx, functionCallEvent, functionResponseEvent) + if got == nil || got.Content == nil || len(got.Content.Parts) != 1 { + t.Fatalf("expected one confirmation part, got %#v", got) + } + if len(got.Content.Parts[0].ThoughtSignature) != 0 { + t.Errorf("ThoughtSignature = %q, want empty", got.Content.Parts[0].ThoughtSignature) + } +} diff --git a/internal/llminternal/googlellm/live_connection.go b/internal/llminternal/googlellm/live_connection.go new file mode 100644 index 000000000..73325e67d --- /dev/null +++ b/internal/llminternal/googlellm/live_connection.go @@ -0,0 +1,294 @@ +// Copyright 2026 Google LLC +// +// Licensed under the Apache License, Version 2.0 (the "License"); +// you may not use this file except in compliance with the License. +// You may obtain a copy of the License at +// +// http://www.apache.org/licenses/LICENSE-2.0 +// +// Unless required by applicable law or agreed to in writing, software +// distributed under the License is distributed on an "AS IS" BASIS, +// WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. +// See the License for the specific language governing permissions and +// limitations under the License. + +package googlellm + +import ( + "context" + "fmt" + "log" + "strings" + + "google.golang.org/genai" + + "google.golang.org/adk/model" +) + +// LiveConnection wraps the underlying GenAI SDK live session. +type LiveConnection struct { + // Using the correct Session type from the GenAI SDK. + sdkSession *genai.Session + + modelName string + backend genai.Backend + inputTranscriptionText string + outputTranscriptionText string + bufferedResponses []*model.LLMResponse +} + +// NewLiveConnection creates a new LiveConnection. +func NewLiveConnection(session *genai.Session, modelName string, backend genai.Backend) *LiveConnection { + return &LiveConnection{ + sdkSession: session, + modelName: modelName, + backend: backend, + } +} + +// SendHistory sends the conversation history to prime the session. +func (c *LiveConnection) SendHistory(ctx context.Context, history []*genai.Content) error { + // TODO: genai seems to be missing initial_history_in_client_content flag + isGemini31 := strings.Contains(c.modelName, "gemini-3.1") + if isGemini31 { + log.Printf("skipping sending history for gemini 3.1\n") + return nil + } + log.Printf("sending preprocessed content %d\n", len(history)) + + var filteredHistory []*genai.Content + for _, content := range history { + if content == nil { + continue + } + var filteredParts []*genai.Part + for _, part := range content.Parts { + if part == nil { + continue + } + if part.InlineData != nil && strings.HasPrefix(part.InlineData.MIMEType, "audio/") { + continue + } + filteredParts = append(filteredParts, part) + } + if len(filteredParts) > 0 { + filteredHistory = append(filteredHistory, &genai.Content{ + Parts: filteredParts, + Role: content.Role, + }) + } + } + log.Printf("sending history: of size %d\n", len(filteredHistory)) + turnComplete := len(filteredHistory) > 0 && filteredHistory[len(filteredHistory)-1].Role == "user" + if len(filteredHistory) > 0 { + err := c.sdkSession.SendClientContent(genai.LiveClientContentInput{ + Turns: filteredHistory, + TurnComplete: &turnComplete, + }) + if err != nil { + return fmt.Errorf("failed to send history: %w", err) + } + } + + return nil +} + +// SendContent sends unary content or function responses to the model. +func (c *LiveConnection) SendContent(ctx context.Context, content *genai.Content) error { + if content == nil || len(content.Parts) == 0 { + return fmt.Errorf("empty content") + } + + if content.Parts[0].FunctionResponse != nil { + var functionResponses []*genai.FunctionResponse + for _, part := range content.Parts { + if part.FunctionResponse != nil { + functionResponses = append(functionResponses, part.FunctionResponse) + } + } + err := c.sdkSession.SendToolResponse(genai.LiveToolResponseInput{ + FunctionResponses: functionResponses, + }) + if err != nil { + return fmt.Errorf("failed to send tool response: %w", err) + } + log.Printf("sending tool response\n") + } else { + isGemini31 := strings.Contains(c.modelName, "gemini-3.1") + isGeminiAPI := c.backend == genai.BackendGeminiAPI + if isGemini31 && isGeminiAPI && len(content.Parts) == 1 && content.Parts[0].Text != "" { + log.Printf("Attempting to send text via SendRealtimeInput\n") + err := c.sdkSession.SendRealtimeInput(genai.LiveRealtimeInput{ + Text: content.Parts[0].Text, + }) + if err != nil { + return fmt.Errorf("failed to send realtime text: %w", err) + } + return nil + } + + turnComplete := true + err := c.sdkSession.SendClientContent(genai.LiveClientContentInput{ + Turns: []*genai.Content{content}, + TurnComplete: &turnComplete, + }) + if err != nil { + return fmt.Errorf("failed to send content: %w", err) + } + } + + return nil +} + +// SendRealtime sends real-time input (audio/video). +func (c *LiveConnection) SendRealtime(ctx context.Context, input any) error { + switch v := input.(type) { + case *genai.Blob: + if v.MIMEType == "" { + // Detect PNG by signature: \x89PNG\r\n\x1a\n + isPNG := len(v.Data) >= 8 && + v.Data[0] == 0x89 && v.Data[1] == 0x50 && v.Data[2] == 0x4E && v.Data[3] == 0x47 && + v.Data[4] == 0x0D && v.Data[5] == 0x0A && v.Data[6] == 0x1A && v.Data[7] == 0x0A + + if isPNG { + v.MIMEType = "image/png" + } else { + v.MIMEType = "audio/pcm" + } + } + + isGemini31 := strings.Contains(c.modelName, "gemini-3.1") + isGeminiAPI := c.backend == genai.BackendGeminiAPI + if isGemini31 && isGeminiAPI { + if strings.HasPrefix(v.MIMEType, "image/") { + return c.sdkSession.SendRealtimeInput(genai.LiveRealtimeInput{ + Video: v, + }) + } + return c.sdkSession.SendRealtimeInput(genai.LiveRealtimeInput{ + Audio: v, + }) + } + + return c.sdkSession.SendRealtimeInput(genai.LiveRealtimeInput{ + Media: v, + }) + case *genai.ActivityStart: + log.Printf("sending activity start\n") + return c.sdkSession.SendRealtimeInput(genai.LiveRealtimeInput{ + ActivityStart: v, + }) + case *genai.ActivityEnd: + log.Printf("sending activity end\n") + return c.sdkSession.SendRealtimeInput(genai.LiveRealtimeInput{ + ActivityEnd: v, + }) + default: + return fmt.Errorf("unsupported real-time input type: %T", input) + } +} + +// Recv receives a response from the live server connection. +func (c *LiveConnection) Recv(ctx context.Context) (*model.LLMResponse, error) { + if len(c.bufferedResponses) > 0 { + resp := c.bufferedResponses[0] + c.bufferedResponses = c.bufferedResponses[1:] + return resp, nil + } + + msg, err := c.sdkSession.Receive() + if err != nil { + return nil, fmt.Errorf("failed to receive message: %w", err) + } + + if msg == nil { + return nil, nil + } + + resp := &model.LLMResponse{} + + if msg.SessionResumptionUpdate != nil { + resp.SessionResumptionHandle = msg.SessionResumptionUpdate.NewHandle + return resp, nil + } + + if msg.ServerContent != nil { + content := msg.ServerContent + if content.ModelTurn != nil { + resp.Content = content.ModelTurn + } + resp.TurnComplete = content.TurnComplete + resp.Interrupted = content.Interrupted + + if content.InputTranscription != nil { + resp.InputTranscription = content.InputTranscription + c.inputTranscriptionText += content.InputTranscription.Text + resp.Partial = true // Mark chunks as partial so they are not saved to session + } + if content.OutputTranscription != nil { + resp.OutputTranscription = content.OutputTranscription + c.outputTranscriptionText += content.OutputTranscription.Text + resp.Partial = true // Mark chunks as partial so they are not saved to session + } + + // Handle transcription finalization on completion signals + if content.TurnComplete || content.Interrupted { + if c.inputTranscriptionText != "" || c.outputTranscriptionText != "" { + if c.inputTranscriptionText != "" { + inputResp := &model.LLMResponse{ + Partial: false, + InputTranscription: &genai.Transcription{ + Text: c.inputTranscriptionText, + Finished: true, + }, + } + c.inputTranscriptionText = "" + c.bufferedResponses = append(c.bufferedResponses, inputResp) + } + if c.outputTranscriptionText != "" { + outputResp := &model.LLMResponse{ + Partial: false, + OutputTranscription: &genai.Transcription{ + Text: c.outputTranscriptionText, + Finished: true, + }, + } + c.outputTranscriptionText = "" + c.bufferedResponses = append(c.bufferedResponses, outputResp) + } + + // Append the current response (which has TurnComplete or Interrupted) to the buffer + // so it is delivered AFTER the transcriptions + c.bufferedResponses = append(c.bufferedResponses, resp) + + // Return the first one from buffer + first := c.bufferedResponses[0] + c.bufferedResponses = c.bufferedResponses[1:] + return first, nil + } + } + } + + if msg.ToolCall != nil { + if resp.Content == nil { + resp.Content = &genai.Content{Role: "model"} + } + for _, call := range msg.ToolCall.FunctionCalls { + if call != nil { + resp.Content.Parts = append(resp.Content.Parts, &genai.Part{ + FunctionCall: call, + }) + } + } + } + + return resp, nil +} + +// Close closes the live server connection. +func (c *LiveConnection) Close() error { + if c.sdkSession != nil { + return c.sdkSession.Close() + } + return nil +} diff --git a/internal/llminternal/handle_function_calls_async_test.go b/internal/llminternal/handle_function_calls_async_test.go index 6137ac89a..5ab985e0b 100644 --- a/internal/llminternal/handle_function_calls_async_test.go +++ b/internal/llminternal/handle_function_calls_async_test.go @@ -38,7 +38,7 @@ type SleepResult struct { Success bool `json:"success"` } -func sleepFunc(ctx tool.Context, input SleepArgs) (SleepResult, error) { +func sleepFunc(ctx agent.ToolContext, input SleepArgs) (SleepResult, error) { time.Sleep(time.Duration(input.DurationMS) * time.Millisecond) return SleepResult{Success: true}, nil } diff --git a/internal/llminternal/liveflow/doc.go b/internal/llminternal/liveflow/doc.go index 93b9f6a52..cc9034ca2 100644 --- a/internal/llminternal/liveflow/doc.go +++ b/internal/llminternal/liveflow/doc.go @@ -33,4 +33,33 @@ // // eventOrError and sendEvent (flow.go) are the shared channel currency // every concern uses to publish events into RunLive's iterator output. +// +// # Relationship to the upstream live engine +// +// Since the google/adk-go v1.5.0 merge, two live engines coexist in this +// repository. Runner.RunLive, the agent.LiveSession API, and the adkrest +// /run_live endpoint drive the UPSTREAM engine (llminternal Flow.RunLive +// plus the googlellm live connection) — not this package. As of v1.5.0 +// that engine has known gaps: +// +// - GoAway is not surfaced as a signal; reconnect eligibility is decided +// by substring-matching error text ("GoAway", "EOF", "1008", ...). +// - Reconnects retry immediately in a loop with no backoff and no +// attempt budget. +// - Each reconnect abandons the previous attempt's unbuffered error +// channel, leaking a blocked reader/sender goroutine per cycle +// (reported upstream: google/adk-go#1152). +// - The preprocessed history is re-sent on every reconnect, even when +// resuming with a session handle the server already has context for. +// - ToolCallCancellation server messages are dropped, so cancelled +// tool calls keep running. +// - Every event is authored as the agent, so input transcriptions +// (user speech) are misattributed to the model in session history. +// +// Runner.RunLiveQueue driving this package is the hulilabs-supported live +// path: it handles each of the above (GoAway-aware reconnects with +// resumption handles, turn-cycle replay suppression, tool cancellation, +// role-correct transcription events). Route new live features and fixes +// here; treat the upstream engine as upstream-owned code that syncs with +// google/adk-go. package liveflow diff --git a/internal/llminternal/liveflow/tools.go b/internal/llminternal/liveflow/tools.go index ee4e52997..c131588a0 100644 --- a/internal/llminternal/liveflow/tools.go +++ b/internal/llminternal/liveflow/tools.go @@ -434,7 +434,7 @@ func (lf *LiveFlow) callToolLive( return map[string]any{"error": err.Error()}, err } - // Rebind invCtx onto ctx so toolinternal.NewToolContext produces a + // Rebind invCtx onto ctx so agent.NewToolContext produces a // tool.Context whose Value() chain reaches the execute_tool span. // Without this, a tool implementation that calls // trace.SpanFromContext(toolCtx) would observe the generate_content @@ -442,7 +442,7 @@ func (lf *LiveFlow) callToolLive( // incorrectly. invCtx = invCtx.WithContext(ctx) - toolCtx := toolinternal.NewToolContext(invCtx, fc.ID, &session.EventActions{StateDelta: make(map[string]any)}, nil) + toolCtx := agent.NewToolContext(invCtx, fc.ID, &session.EventActions{StateDelta: make(map[string]any)}, nil) wrappedCtx := &cancelableToolContext{Context: toolCtx, cancelCtx: ctx} response, err := lf.runToolWithCallbacks(wrappedCtx, invCtx, funcTool, fc.Args) diff --git a/internal/llminternal/outputschema_processor.go b/internal/llminternal/outputschema_processor.go index e64847532..bc87ec0f7 100644 --- a/internal/llminternal/outputschema_processor.go +++ b/internal/llminternal/outputschema_processor.go @@ -27,7 +27,6 @@ import ( "google.golang.org/adk/internal/utils" "google.golang.org/adk/model" "google.golang.org/adk/session" - "google.golang.org/adk/tool" ) const ( @@ -67,7 +66,7 @@ func outputSchemaRequestProcessor(ctx agent.InvocationContext, req *model.LLMReq // createFinalModelResponseEvent creates a final model response event from set_model_response JSON. func createFinalModelResponseEvent(invocationContext agent.InvocationContext, response string) *session.Event { // Create a proper model response event - finalEvent := session.NewEvent(invocationContext.InvocationID()) + finalEvent := session.NewEventWithContext(invocationContext, invocationContext.InvocationID()) finalEvent.Author = invocationContext.Agent().Name() finalEvent.Branch = invocationContext.Branch() finalEvent.Content = &genai.Content{ @@ -129,7 +128,7 @@ func (t *setModelResponseTool) Declaration() *genai.FunctionDeclaration { } } -func (t *setModelResponseTool) Run(ctx tool.Context, args any) (map[string]any, error) { +func (t *setModelResponseTool) Run(ctx agent.ToolContext, args any) (map[string]any, error) { m, ok := args.(map[string]any) if !ok { return nil, fmt.Errorf("unexpected args type for set_model_response: %T", args) diff --git a/internal/llminternal/outputschema_processor_test.go b/internal/llminternal/outputschema_processor_test.go index 6d0d8c29c..aeb97e665 100644 --- a/internal/llminternal/outputschema_processor_test.go +++ b/internal/llminternal/outputschema_processor_test.go @@ -25,7 +25,6 @@ import ( "google.golang.org/adk/agent" icontext "google.golang.org/adk/internal/context" - "google.golang.org/adk/internal/toolinternal" "google.golang.org/adk/internal/utils" "google.golang.org/adk/model" "google.golang.org/adk/session" @@ -42,7 +41,7 @@ func (m *mockTool) IsLongRunning() bool { return false } func (m *mockTool) Declaration() *genai.FunctionDeclaration { return &genai.FunctionDeclaration{Name: m.name} } -func (m *mockTool) Run(ctx tool.Context, args any) (map[string]any, error) { return nil, nil } +func (m *mockTool) Run(ctx agent.ToolContext, args any) (map[string]any, error) { return nil, nil } type mockLLM struct { model.LLM @@ -315,7 +314,7 @@ func TestSetModelResponseTool(t *testing.T) { t.Run("RunSuccess", func(t *testing.T) { invCtx := icontext.NewInvocationContext(context.Background(), icontext.InvocationContextParams{}) - toolCtx := toolinternal.NewToolContext(invCtx, "", nil, nil) + toolCtx := agent.NewToolContext(invCtx, "", nil, nil) input := map[string]any{"count": 10.0} // JSON numbers often come as float64 got, err := toolInstance.Run(toolCtx, input) @@ -329,7 +328,7 @@ func TestSetModelResponseTool(t *testing.T) { t.Run("RunValidationFailure_Type", func(t *testing.T) { invCtx := icontext.NewInvocationContext(context.Background(), icontext.InvocationContextParams{}) - toolCtx := toolinternal.NewToolContext(invCtx, "", nil, nil) + toolCtx := agent.NewToolContext(invCtx, "", nil, nil) input := map[string]any{"count": "not a number"} _, err := toolInstance.Run(toolCtx, input) @@ -340,7 +339,7 @@ func TestSetModelResponseTool(t *testing.T) { t.Run("RunValidationFailure_MissingRequired", func(t *testing.T) { invCtx := icontext.NewInvocationContext(context.Background(), icontext.InvocationContextParams{}) - toolCtx := toolinternal.NewToolContext(invCtx, "", nil, nil) + toolCtx := agent.NewToolContext(invCtx, "", nil, nil) input := map[string]any{"other": 123} _, err := toolInstance.Run(toolCtx, input) diff --git a/internal/llminternal/parallel_function_call_hitl_test.go b/internal/llminternal/parallel_function_call_hitl_test.go new file mode 100644 index 000000000..00fbc47c6 --- /dev/null +++ b/internal/llminternal/parallel_function_call_hitl_test.go @@ -0,0 +1,464 @@ +// Copyright 2026 Google LLC +// +// Licensed under the Apache License, Version 2.0 (the "License"); +// you may not use this file except in compliance with the License. +// You may obtain a copy of the License at +// +// http://www.apache.org/licenses/LICENSE-2.0 +// +// Unless required by applicable law or agreed to in writing, software +// distributed under the License is distributed on an "AS IS" BASIS, +// WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. +// See the License for the specific language governing permissions and +// limitations under the License. + +package llminternal_test + +import ( + "context" + "encoding/json" + "iter" + "testing" + + "github.com/google/go-cmp/cmp" + "github.com/google/go-cmp/cmp/cmpopts" + "google.golang.org/genai" + + "google.golang.org/adk/agent" + "google.golang.org/adk/agent/llmagent" + "google.golang.org/adk/model" + "google.golang.org/adk/runner" + "google.golang.org/adk/session" + "google.golang.org/adk/tool" + "google.golang.org/adk/tool/functiontool" + "google.golang.org/adk/tool/toolconfirmation" +) + +type SecureActionArgs struct { + ActionName string `json:"action_name"` +} + +type SecureActionResult struct { + Executed bool `json:"executed"` +} + +func secureActionFunc(ctx agent.ToolContext, input SecureActionArgs) (SecureActionResult, error) { + return SecureActionResult{Executed: true}, nil +} + +type hitlMockModel struct { + model.LLM + Calls int +} + +func (m *hitlMockModel) Name() string { + return "hitl-mock-model" +} + +func (m *hitlMockModel) GenerateContent(ctx context.Context, req *model.LLMRequest, useStream bool) iter.Seq2[*model.LLMResponse, error] { + return func(yield func(*model.LLMResponse, error) bool) { + m.Calls++ + if m.Calls > 1 { + // Turn 2 model execution (after user confirms) + yield(&model.LLMResponse{ + Content: &genai.Content{ + Parts: []*genai.Part{ + genai.NewPartFromText("Successfully processed your request after confirmation."), + }, + Role: "model", + }, + Partial: false, + }, nil) + return + } + + // Turn 1: yields 2 parallel function calls requiring confirmation + yield(&model.LLMResponse{ + Content: &genai.Content{ + Parts: []*genai.Part{ + { + FunctionCall: &genai.FunctionCall{ + ID: "secure_call_1", + Name: "secure_action", + Args: map[string]any{"action_name": "deploy_prod"}, + }, + }, + { + FunctionCall: &genai.FunctionCall{ + ID: "secure_call_2", + Name: "secure_action", + Args: map[string]any{"action_name": "delete_db"}, + }, + }, + }, + Role: "model", + }, + Partial: false, + }, nil) + } +} + +// TestParallelFunctionCallsWithHITL verifies the end-to-end coordination of multiple parallel +// tool executions that require Human-in-the-Loop (HITL) confirmation. +// +// The test simulates a two-turn interaction: +// 1. Turn 1: The model initiates two parallel tool calls to a sensitive tool. Because no user +// confirmation is yet present, both tools trigger RequestConfirmation. This sets the +// SkipSummarization flag, which halts LLM generation immediately after tool responses are returned. +// The runner yields: +// a) A model response event containing the two original parallel function calls. +// b) A merged function response event containing the placeholder (unexecuted) tool responses. +// c) An aggregated tool confirmation event containing two adk_request_confirmation wrapper calls. +// 2. Turn 2: The client simulates user confirmation by returning confirmation responses matching +// the unique wrapper call IDs. The RequestConfirmationRequestProcessor detects these, matches +// them back to the original tools, and concurrently executes the sensitive tools in parallel +// with the confirmed flag enabled. Finally, the model is called to generate a summary. +func TestParallelFunctionCallsWithHITL(t *testing.T) { + secureTool, err := functiontool.New(functiontool.Config{ + Name: "secure_action", + Description: "performs a sensitive/secure action", + RequireConfirmation: true, + }, secureActionFunc) + if err != nil { + t.Fatal(err) + } + + mockModel := &hitlMockModel{} + + a, err := llmagent.New(llmagent.Config{ + Name: "hitl_tester", + Description: "HITL tester agent", + Instruction: "You are a secure helper.", + Model: mockModel, + Tools: []tool.Tool{ + secureTool, + }, + }) + if err != nil { + t.Fatal(err) + } + + sessionService := session.InMemoryService() + _, err = sessionService.Create(t.Context(), &session.CreateRequest{ + AppName: "testApp", + UserID: "testUser", + SessionID: "testSession", + }) + if err != nil { + t.Fatal(err) + } + + r, err := runner.New(runner.Config{ + Agent: a, + SessionService: sessionService, + AppName: "testApp", + }) + if err != nil { + t.Fatal(err) + } + + // --- TURN 1: Requesting parallel tools requiring confirmation --- + it := r.Run(t.Context(), "testUser", "testSession", &genai.Content{ + Parts: []*genai.Part{ + genai.NewPartFromText("Please run prod deploy and db delete"), + }, + Role: "user", + }, agent.RunConfig{StreamingMode: agent.StreamingModeSSE}) + + var turn1Events []*session.Event + for ev, err := range it { + if err != nil { + t.Fatal(err) + } + turn1Events = append(turn1Events, ev) + } + + // Expecting exactly 3 events: + // 1. ModelResponseEvent (yielding the two function calls secure_call_1 and secure_call_2) + // 2. Merged function response event (returning placeholder executed: false responses) + // 3. ToolConfirmationEvent (yielding two adk_request_confirmation calls) + if len(turn1Events) != 3 { + t.Fatalf("Expected 3 events in Turn 1, got %d", len(turn1Events)) + } + + // Verify that the model event contains our parallel calls + modelEvent := turn1Events[0] + if len(modelEvent.Content.Parts) != 2 { + t.Errorf("Expected model event to contain 2 parts (function calls), got %d", len(modelEvent.Content.Parts)) + } + + // Verify that the merged function response event contains the placeholder unexecuted responses + placeholderRespEvent := turn1Events[1] + if len(placeholderRespEvent.Content.Parts) != 2 { + t.Errorf("Expected placeholder function response event to contain 2 parts, got %d", len(placeholderRespEvent.Content.Parts)) + } + + expectedPlaceholders := []*genai.Part{ + { + FunctionResponse: &genai.FunctionResponse{ + Name: "secure_action", + ID: "secure_call_1", + Response: map[string]any{"error": `error tool "secure_action" requires confirmation, please approve or reject`}, + }, + }, + { + FunctionResponse: &genai.FunctionResponse{ + Name: "secure_action", + ID: "secure_call_2", + Response: map[string]any{"error": `error tool "secure_action" requires confirmation, please approve or reject`}, + }, + }, + } + + if diff := cmp.Diff(expectedPlaceholders, placeholderRespEvent.Content.Parts); diff != "" { + t.Errorf("Mismatch in placeholder tool responses (-want +got):\n%s", diff) + } + + // Verify that the tool confirmation event has the confirmation wrapper calls + confirmEvent := turn1Events[2] + if len(confirmEvent.Content.Parts) != 2 { + t.Errorf("Expected tool confirmation event to contain 2 wrapper function calls, got %d", len(confirmEvent.Content.Parts)) + } + + var confirmCallID1, confirmCallID2 string + for _, p := range confirmEvent.Content.Parts { + if p.FunctionCall == nil || p.FunctionCall.Name != toolconfirmation.FunctionCallName { + t.Errorf("Expected function call name %s, got %v", toolconfirmation.FunctionCallName, p.FunctionCall) + continue + } + origCall, err := toolconfirmation.OriginalCallFrom(p.FunctionCall) + if err != nil { + t.Fatalf("Failed to extract original call: %v", err) + } + switch origCall.ID { + case "secure_call_1": + confirmCallID1 = p.FunctionCall.ID + case "secure_call_2": + confirmCallID2 = p.FunctionCall.ID + } + } + + if confirmCallID1 == "" || confirmCallID2 == "" { + t.Fatalf("Failed to retrieve both confirmation function call IDs from turn 1 events") + } + + // --- TURN 2: User confirms the actions --- + // Build user's response to confirmation requests. + userConfirmation := toolconfirmation.ToolConfirmation{Confirmed: true} + userConfirmationJSON, _ := json.Marshal(userConfirmation) + userConfirmationResponse := map[string]any{ + "response": string(userConfirmationJSON), + } + + // Run runner again, passing the user confirmation response content directly as the message + it2 := r.Run(t.Context(), "testUser", "testSession", &genai.Content{ + Role: "user", + Parts: []*genai.Part{ + { + FunctionResponse: &genai.FunctionResponse{ + Name: toolconfirmation.FunctionCallName, + ID: confirmCallID1, + Response: userConfirmationResponse, + }, + }, + { + FunctionResponse: &genai.FunctionResponse{ + Name: toolconfirmation.FunctionCallName, + ID: confirmCallID2, + Response: userConfirmationResponse, + }, + }, + }, + }, agent.RunConfig{StreamingMode: agent.StreamingModeSSE}) + + var turn2Events []*session.Event + for ev, err := range it2 { + if err != nil { + t.Fatal(err) + } + turn2Events = append(turn2Events, ev) + } + + // Turn 2 Events Expectation: + // 1. Function response event containing executed: true for both calls + // 2. Final summarizing model response ("Successfully processed your request after confirmation.") + if len(turn2Events) != 2 { + t.Fatalf("Expected 2 events in Turn 2, got %d", len(turn2Events)) + } + + funcRespEvent := turn2Events[0] + if len(funcRespEvent.Content.Parts) != 2 { + t.Errorf("Expected 2 function responses in Turn 2, got %d", len(funcRespEvent.Content.Parts)) + } + + ignoreFields := []cmp.Option{ + cmpopts.IgnoreFields(genai.FunctionResponse{}, "ID"), + } + + expectedResponses := []*genai.Part{ + { + FunctionResponse: &genai.FunctionResponse{ + Name: "secure_action", + Response: map[string]any{"executed": true}, + }, + }, + { + FunctionResponse: &genai.FunctionResponse{ + Name: "secure_action", + Response: map[string]any{"executed": true}, + }, + }, + } + + if diff := cmp.Diff(expectedResponses, funcRespEvent.Content.Parts, ignoreFields...); diff != "" { + t.Errorf("Mismatch in executed tool responses (-want +got):\n%s", diff) + } + + finalModelEvent := turn2Events[1] + if finalModelEvent.Content.Parts[0].Text != "Successfully processed your request after confirmation." { + t.Errorf("Expected summarizing model event, got: %v", finalModelEvent.Content.Parts[0]) + } +} + +// TestParallelFunctionCallsWithPartialHITL verifies the behavior when multiple parallel +// tool executions are requested, but only one of them is confirmed by the user while the +// other is rejected/denied. +func TestParallelFunctionCallsWithPartialHITL(t *testing.T) { + secureTool, err := functiontool.New(functiontool.Config{ + Name: "secure_action", + Description: "performs a sensitive/secure action", + RequireConfirmation: true, + }, secureActionFunc) + if err != nil { + t.Fatal(err) + } + + mockModel := &hitlMockModel{} + + a, err := llmagent.New(llmagent.Config{ + Name: "hitl_tester", + Description: "HITL tester agent", + Instruction: "You are a secure helper.", + Model: mockModel, + Tools: []tool.Tool{ + secureTool, + }, + }) + if err != nil { + t.Fatal(err) + } + + sessionService := session.InMemoryService() + _, err = sessionService.Create(t.Context(), &session.CreateRequest{ + AppName: "testApp", + UserID: "testUser", + SessionID: "testSession", + }) + if err != nil { + t.Fatal(err) + } + + r, err := runner.New(runner.Config{ + Agent: a, + SessionService: sessionService, + AppName: "testApp", + }) + if err != nil { + t.Fatal(err) + } + + // --- TURN 1: Requesting parallel tools requiring confirmation --- + it := r.Run(t.Context(), "testUser", "testSession", &genai.Content{ + Parts: []*genai.Part{ + genai.NewPartFromText("Please run prod deploy and db delete"), + }, + Role: "user", + }, agent.RunConfig{StreamingMode: agent.StreamingModeSSE}) + + var turn1Events []*session.Event + for ev, err := range it { + if err != nil { + t.Fatal(err) + } + turn1Events = append(turn1Events, ev) + } + + if len(turn1Events) != 3 { + t.Fatalf("Expected 3 events in Turn 1, got %d", len(turn1Events)) + } + + confirmEvent := turn1Events[2] + var confirmCallID1, confirmCallID2 string + for _, p := range confirmEvent.Content.Parts { + origCall, err := toolconfirmation.OriginalCallFrom(p.FunctionCall) + if err != nil { + t.Fatalf("Failed to extract original call: %v", err) + } + switch origCall.ID { + case "secure_call_1": + confirmCallID1 = p.FunctionCall.ID + case "secure_call_2": + confirmCallID2 = p.FunctionCall.ID + } + } + + if confirmCallID1 == "" || confirmCallID2 == "" { + t.Fatalf("Failed to retrieve both confirmation function call IDs from turn 1 events") + } + + // --- TURN 2: User partially confirms the actions --- + // Only confirm secure_call_2 (secure_call_1 remains pending/unconfirmed in this turn). + confirmedJSON, _ := json.Marshal(toolconfirmation.ToolConfirmation{Confirmed: true}) + + it2 := r.Run(t.Context(), "testUser", "testSession", &genai.Content{ + Role: "user", + Parts: []*genai.Part{ + { + FunctionResponse: &genai.FunctionResponse{ + Name: toolconfirmation.FunctionCallName, + ID: confirmCallID2, + Response: map[string]any{"response": string(confirmedJSON)}, + }, + }, + }, + }, agent.RunConfig{StreamingMode: agent.StreamingModeSSE}) + + var turn2Events []*session.Event + for ev, err := range it2 { + if err != nil { + t.Fatal(err) + } + turn2Events = append(turn2Events, ev) + } + + if len(turn2Events) != 2 { + t.Fatalf("Expected 2 events in Turn 2, got %d", len(turn2Events)) + } + + funcRespEvent := turn2Events[0] + if len(funcRespEvent.Content.Parts) != 1 { + t.Errorf("Expected 1 function response in Turn 2, got %d", len(funcRespEvent.Content.Parts)) + } + + ignoreFields := []cmp.Option{ + cmpopts.IgnoreFields(genai.FunctionResponse{}, "ID"), + } + + expectedResponses := []*genai.Part{ + { + FunctionResponse: &genai.FunctionResponse{ + Name: "secure_action", + Response: map[string]any{"executed": true}, + }, + }, + } + + if diff := cmp.Diff(expectedResponses, funcRespEvent.Content.Parts, ignoreFields...); diff != "" { + t.Errorf("Mismatch in executed tool responses (-want +got):\n%s", diff) + } + + finalModelEvent := turn2Events[1] + if finalModelEvent.Content.Parts[0].Text != "Successfully processed your request after confirmation." { + t.Errorf("Expected summarizing model event, got: %v", finalModelEvent.Content.Parts[0]) + } +} diff --git a/internal/llminternal/parallel_function_call_test.go b/internal/llminternal/parallel_function_call_test.go index 80e1a767e..c115af3f4 100644 --- a/internal/llminternal/parallel_function_call_test.go +++ b/internal/llminternal/parallel_function_call_test.go @@ -44,7 +44,7 @@ type SumResult struct { Sum int `json:"sum"` // the sum of two integers } -func sumFunc(ctx tool.Context, input SumArgs) (SumResult, error) { +func sumFunc(ctx agent.ToolContext, input SumArgs) (SumResult, error) { return SumResult{Sum: input.A + input.B}, nil } diff --git a/internal/llminternal/request_confirmation_processor.go b/internal/llminternal/request_confirmation_processor.go index 16aab5304..b87fad39f 100644 --- a/internal/llminternal/request_confirmation_processor.go +++ b/internal/llminternal/request_confirmation_processor.go @@ -163,7 +163,7 @@ func RequestConfirmationRequestProcessor(ctx agent.InvocationContext, req *model ev, err := f.handleFunctionCalls(ctx, toolsmap, &model.LLMResponse{ Content: &genai.Content{Parts: parts, Role: genai.RoleUser}, - }, toolsToResumeConfirmation) + }, toolsToResumeConfirmation, nil) if !yield(ev, err) { return } diff --git a/internal/llminternal/request_confirmation_processor_test.go b/internal/llminternal/request_confirmation_processor_test.go index 3a881d23c..bae287ad2 100644 --- a/internal/llminternal/request_confirmation_processor_test.go +++ b/internal/llminternal/request_confirmation_processor_test.go @@ -51,7 +51,7 @@ func (m *mockTool) Declaration() *genai.FunctionDeclaration { return &genai.FunctionDeclaration{Name: m.name} } -func (m *mockTool) Run(ctx tool.Context, args any) (map[string]any, error) { +func (m *mockTool) Run(ctx agent.ToolContext, args any) (map[string]any, error) { if ctx.ToolConfirmation() == nil || !ctx.ToolConfirmation().Confirmed { return map[string]any{"error": string("Tool execution not confirmed")}, nil } diff --git a/internal/llminternal/stream_aggregator.go b/internal/llminternal/stream_aggregator.go index ac7b2d0af..a105d5ac6 100644 --- a/internal/llminternal/stream_aggregator.go +++ b/internal/llminternal/stream_aggregator.go @@ -16,7 +16,6 @@ package llminternal import ( "context" - "fmt" "iter" "maps" "reflect" @@ -37,6 +36,8 @@ type streamingResponseAggregator struct { citationMetadata *genai.CitationMetadata response *model.LLMResponse + currentThoughtSignature []byte + sequence []*genai.Part currentTextBuffer string currentTextIsThought bool @@ -57,14 +58,11 @@ func NewStreamingResponseAggregator() *streamingResponseAggregator { // also yielding an aggregated response if the GenerateContentResponse has zero parts or is audio data func (s *streamingResponseAggregator) ProcessResponse(ctx context.Context, genResp *genai.GenerateContentResponse) iter.Seq2[*model.LLMResponse, error] { return func(yield func(*model.LLMResponse, error) bool) { - if len(genResp.Candidates) == 0 { - // shouldn't happen? - yield(nil, fmt.Errorf("empty response")) - return - } - candidate := genResp.Candidates[0] resp := converters.Genai2LLMResponse(genResp) - resp.TurnComplete = candidate.FinishReason != "" + if len(genResp.Candidates) > 0 { + candidate := genResp.Candidates[0] + resp.TurnComplete = candidate.FinishReason != "" + } // Aggregate the response and check if an intermediate event to yield was created if aggrResp := s.aggregateResponse(resp); aggrResp != nil { if !yield(aggrResp, nil) { @@ -102,6 +100,9 @@ func (s *streamingResponseAggregator) aggregateResponse(llmResponse *model.LLMRe if reflect.ValueOf(*part).IsZero() { continue } + if len(part.ThoughtSignature) > 0 { + s.currentThoughtSignature = part.ThoughtSignature + } if part.Text != "" { if s.currentTextBuffer != "" && part.Thought != s.currentTextIsThought { s.flushTextBufferToSequence() @@ -135,6 +136,10 @@ func (s *streamingResponseAggregator) processFunctionCallPart(part *genai.Part) } else { if part.FunctionCall.Name != "" { s.flushTextBufferToSequence() + if part.ThoughtSignature == nil && s.currentThoughtSignature != nil { + part.ThoughtSignature = s.currentThoughtSignature + } + s.currentThoughtSignature = nil s.sequence = append(s.sequence, part) } } diff --git a/internal/llminternal/stream_aggregator_test.go b/internal/llminternal/stream_aggregator_test.go index 0e18153ee..3889b89f0 100644 --- a/internal/llminternal/stream_aggregator_test.go +++ b/internal/llminternal/stream_aggregator_test.go @@ -148,6 +148,160 @@ func TestProgressiveSSEStreamingFunctionCallArguments(t *testing.T) { } } +func TestThoughtSignaturePropagationToFirstFunctionCallSeparateResponses(t *testing.T) { + aggregator := llminternal.NewStreamingResponseAggregator() + ctx := t.Context() + + testThoughtSignature := []byte("test_signature_123") + + chunk1 := &genai.GenerateContentResponse{ + Candidates: []*genai.Candidate{ + { + Content: &genai.Content{ + Role: "model", + Parts: []*genai.Part{ + { + // Emulate an executable code part carrying thought signature + ThoughtSignature: testThoughtSignature, + }, + }, + }, + }, + }, + } + + chunk2 := &genai.GenerateContentResponse{ + Candidates: []*genai.Candidate{ + { + Content: &genai.Content{ + Role: "model", + Parts: []*genai.Part{ + { + FunctionCall: &genai.FunctionCall{ + Name: "print", + Args: map[string]any{"message": "hello"}, + }, + }, + { + FunctionCall: &genai.FunctionCall{ + Name: "print", + Args: map[string]any{"message": "world"}, + }, + }, + }, + }, + FinishReason: genai.FinishReasonStop, + }, + }, + } + + for _, chunk := range []*genai.GenerateContentResponse{chunk1, chunk2} { + for _, err := range aggregator.ProcessResponse(ctx, chunk) { + if err != nil { + t.Fatalf("unexpected error: %v", err) + } + } + } + + finalResponse := aggregator.Close() + if finalResponse == nil { + t.Fatal("expected final response") + } + + parts := finalResponse.Content.Parts + if len(parts) != 3 { + t.Fatalf("expected 3 parts, got %d", len(parts)) + } + + // The first function call should get the propagated thought signature + fcPart1 := parts[1] + if fcPart1.FunctionCall == nil || fcPart1.FunctionCall.Name != "print" { + t.Fatal("expected first function call to be print") + } + if string(fcPart1.ThoughtSignature) != string(testThoughtSignature) { + t.Errorf("expected first function call to have thought signature %s, got %s", string(testThoughtSignature), string(fcPart1.ThoughtSignature)) + } + + // The second function call should NOT get the signature (as intended by user) + fcPart2 := parts[2] + if fcPart2.FunctionCall == nil || fcPart2.FunctionCall.Name != "print" { + t.Fatal("expected second function call to be print") + } + if len(fcPart2.ThoughtSignature) > 0 { + t.Errorf("expected second function call to have no thought signature, but got %s", string(fcPart2.ThoughtSignature)) + } +} + +func TestThoughtSignaturePropagationToFirstFunctionCallSingleResponse(t *testing.T) { + aggregator := llminternal.NewStreamingResponseAggregator() + ctx := t.Context() + + testThoughtSignature := []byte("test_signature_123") + + chunk := &genai.GenerateContentResponse{ + Candidates: []*genai.Candidate{ + { + Content: &genai.Content{ + Role: "model", + Parts: []*genai.Part{ + { + // Emulate an executable code part carrying thought signature + ThoughtSignature: testThoughtSignature, + }, + { + FunctionCall: &genai.FunctionCall{ + Name: "print", + Args: map[string]any{"message": "hello"}, + }, + }, + { + FunctionCall: &genai.FunctionCall{ + Name: "print", + Args: map[string]any{"message": "world"}, + }, + }, + }, + }, + FinishReason: genai.FinishReasonStop, + }, + }, + } + + for _, err := range aggregator.ProcessResponse(ctx, chunk) { + if err != nil { + t.Fatalf("unexpected error: %v", err) + } + } + + finalResponse := aggregator.Close() + if finalResponse == nil { + t.Fatal("expected final response") + } + + parts := finalResponse.Content.Parts + if len(parts) != 3 { + t.Fatalf("expected 3 parts, got %d", len(parts)) + } + + // The first function call should get the propagated thought signature + fcPart1 := parts[1] + if fcPart1.FunctionCall == nil || fcPart1.FunctionCall.Name != "print" { + t.Fatal("expected first function call to be print") + } + if string(fcPart1.ThoughtSignature) != string(testThoughtSignature) { + t.Errorf("expected first function call to have thought signature %s, got %s", string(testThoughtSignature), string(fcPart1.ThoughtSignature)) + } + + // The second function call should NOT get the signature (as intended by user) + fcPart2 := parts[2] + if fcPart2.FunctionCall == nil || fcPart2.FunctionCall.Name != "print" { + t.Fatal("expected second function call to be print") + } + if len(fcPart2.ThoughtSignature) > 0 { + t.Errorf("expected second function call to have no thought signature, but got %s", string(fcPart2.ThoughtSignature)) + } +} + func TestProgressiveSSEPreservesThoughtSignature(t *testing.T) { aggregator := llminternal.NewStreamingResponseAggregator() ctx := t.Context() @@ -475,7 +629,7 @@ type GetWeatherArgs struct { Location string `json:"location"` } -func getWeather(ctx tool.Context, args GetWeatherArgs) (map[string]any, error) { +func getWeather(ctx agent.ToolContext, args GetWeatherArgs) (map[string]any, error) { return map[string]any{ "temperature": 22, "condition": "sunny", @@ -791,7 +945,7 @@ func TestPartialFunctionCallsNotExecutedInNoneStreamingMode(t *testing.T) { CallID string `json:"call_id"` } - trackExecution := func(ctx tool.Context, args TrackExecutionArgs) (string, error) { + trackExecution := func(ctx agent.ToolContext, args TrackExecutionArgs) (string, error) { executionLog = append(executionLog, args.CallID) return "Executed: " + args.CallID, nil } @@ -847,6 +1001,109 @@ func TestPartialFunctionCallsNotExecutedInNoneStreamingMode(t *testing.T) { } } +func TestMetadataVertexAISSEStream(t *testing.T) { + aggregator := llminternal.NewStreamingResponseAggregator() + ctx := t.Context() + + emptyTextChunk := &genai.GenerateContentResponse{ + Candidates: []*genai.Candidate{ + {Content: genai.NewContentFromText("", "model")}, + }, + UsageMetadata: &genai.GenerateContentResponseUsageMetadata{}, + } + + // Simulate the leading metadata-only SSE chunk that Vertex AI emits for + // gemini-3-flash-preview + googleSearch grounding: no Candidates at all. + metadataChunk := &genai.GenerateContentResponse{ + // Candidates is intentionally nil / empty. + UsageMetadata: &genai.GenerateContentResponseUsageMetadata{}, + } + + // The real content arrives in the next chunk. + contentChunk := &genai.GenerateContentResponse{ + Candidates: []*genai.Candidate{ + { + Content: &genai.Content{ + Role: "model", + Parts: []*genai.Part{{Text: "Here are some movie recommendations."}}, + }, + FinishReason: genai.FinishReasonStop, + }, + }, + } + + stream := []*genai.GenerateContentResponse{emptyTextChunk, metadataChunk, metadataChunk, emptyTextChunk, metadataChunk, metadataChunk, contentChunk} + + for _, chunk := range stream { + for resp, err := range aggregator.ProcessResponse(ctx, chunk) { + if err != nil { + t.Fatalf("unexpected error processing chunk: %v", err) + } + _ = resp + } + } + + finalResponse := aggregator.Close() + if finalResponse == nil { + t.Fatal("expected a final aggregated response, got nil") + } + + if len(finalResponse.Content.Parts) == 0 { + t.Fatal("expected content parts in the final response, got none") + } + + if finalResponse.Content.Parts[0].Text != "Here are some movie recommendations." { + t.Errorf("unexpected text in final response: %q", finalResponse.Content.Parts[0].Text) + } +} + +func TestMetadataOnlyChunkDoesNotAbortStream(t *testing.T) { + aggregator := llminternal.NewStreamingResponseAggregator() + ctx := t.Context() + + // Simulate the leading metadata-only SSE chunk that Vertex AI emits for + // gemini-3-flash-preview + googleSearch grounding: no Candidates at all. + metadataChunk := &genai.GenerateContentResponse{ + // Candidates is intentionally nil / empty. + UsageMetadata: &genai.GenerateContentResponseUsageMetadata{}, + } + + // The real content arrives in the next chunk. + contentChunk := &genai.GenerateContentResponse{ + Candidates: []*genai.Candidate{ + { + Content: &genai.Content{ + Role: "model", + Parts: []*genai.Part{{Text: "Here are some movie recommendations."}}, + }, + FinishReason: genai.FinishReasonStop, + }, + }, + } + + for _, chunk := range []*genai.GenerateContentResponse{metadataChunk, contentChunk} { + for resp, err := range aggregator.ProcessResponse(ctx, chunk) { + if err != nil { + t.Fatalf("unexpected error processing chunk: %v", err) + } + _ = resp + } + } + + finalResponse := aggregator.Close() + if finalResponse == nil { + t.Fatal("expected a final aggregated response, got nil") + } + + if len(finalResponse.Content.Parts) == 0 { + t.Fatal("expected content parts in the final response, got none") + } + + if finalResponse.Content.Parts[0].Text != "Here are some movie recommendations." { + t.Errorf("unexpected text in final response: %q", finalResponse.Content.Parts[0].Text) + } +} + func TestFinishReasonUnexpectedToolCallPreservesErrorCode(t *testing.T) { aggregator := llminternal.NewStreamingResponseAggregator() ctx := t.Context() diff --git a/internal/llminternal/streaming_tool_test.go b/internal/llminternal/streaming_tool_test.go new file mode 100644 index 000000000..763dc7e4f --- /dev/null +++ b/internal/llminternal/streaming_tool_test.go @@ -0,0 +1,322 @@ +// Copyright 2026 Google LLC +// +// Licensed under the Apache License, Version 2.0 (the "License"); +// you may not use this file except in compliance with the License. +// You may obtain a copy of the License at +// +// http://www.apache.org/licenses/LICENSE-2.0 +// +// Unless required by applicable law or agreed to in writing, software +// distributed under the License is distributed on an "AS IS" BASIS, +// WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. +// See the License for the specific language governing permissions and +// limitations under the License. + +package llminternal + +import ( + "fmt" + "iter" + "sync" + "testing" + "time" + + "github.com/google/go-cmp/cmp" + "google.golang.org/genai" + + "google.golang.org/adk/agent" + icontext "google.golang.org/adk/internal/context" + "google.golang.org/adk/model" + "google.golang.org/adk/tool" + "google.golang.org/adk/tool/functiontool" +) + +type mockLiveSession struct { + sendFunc func(agent.LiveRequest) error +} + +func (m *mockLiveSession) Send(req agent.LiveRequest) error { + if m.sendFunc != nil { + return m.sendFunc(req) + } + return nil +} + +func (m *mockLiveSession) Close() error { return nil } + +func TestHandleFunctionCalls_Streaming(t *testing.T) { + type Args struct { + Count int `json:"count"` + } + + handler := func(ctx agent.ToolContext, args Args) iter.Seq2[string, error] { + return func(yield func(string, error) bool) { + for i := 0; i < args.Count; i++ { + if !yield(fmt.Sprintf("chunk %d", i), nil) { + return + } + } + } + } + + streamTool, err := functiontool.NewStreaming(functiontool.Config{ + Name: "test_stream", + Description: "streams chunks", + }, handler) + if err != nil { + t.Fatal(err) + } + + toolsDict := map[string]tool.Tool{ + "test_stream": streamTool, + } + + resp := &model.LLMResponse{ + Content: &genai.Content{ + Parts: []*genai.Part{ + { + FunctionCall: &genai.FunctionCall{ + ID: "call_1", + Name: "test_stream", + Args: map[string]any{"count": 3}, + }, + }, + }, + Role: "model", + }, + } + + t.Run("Live Mode (Streaming)", func(t *testing.T) { + var receivedChunks []string + var mu sync.Mutex + var wg sync.WaitGroup + wg.Add(3) // We expect 3 chunks + + mockSess := &mockLiveSession{ + sendFunc: func(req agent.LiveRequest) error { + mu.Lock() + defer mu.Unlock() + if req.Content != nil && len(req.Content.Parts) > 0 { + receivedChunks = append(receivedChunks, req.Content.Parts[0].Text) + } + wg.Done() + return nil + }, + } + + invCtx := icontext.NewInvocationContext(t.Context(), icontext.InvocationContextParams{ + InvocationID: "inv_1", + Agent: &mockAgent{name: "agent_1"}, + }) + + flow := &Flow{ + Tools: []tool.Tool{streamTool}, + } + + mergedEvent, err := flow.handleFunctionCalls(invCtx, toolsDict, resp, nil, mockSess) + if err != nil { + t.Fatal(err) + } + + // Verify immediate response is pending + if mergedEvent == nil { + t.Fatal("expected non-nil mergedEvent") + } + if len(mergedEvent.LLMResponse.Content.Parts) != 1 { + t.Fatalf("expected 1 part, got %d", len(mergedEvent.LLMResponse.Content.Parts)) + } + respPart := mergedEvent.LLMResponse.Content.Parts[0].FunctionResponse + if respPart == nil { + t.Fatal("expected FunctionResponse part") + } + status, ok := respPart.Response["status"].(string) + if !ok || status != "The function is running asynchronously and the results are pending." { + t.Errorf("unexpected status: %v", respPart.Response["status"]) + } + + // Wait for background streaming to complete + wg.Wait() + + wantChunks := []string{ + "Function test_stream returned: chunk 0", + "Function test_stream returned: chunk 1", + "Function test_stream returned: chunk 2", + } + if diff := cmp.Diff(wantChunks, receivedChunks); diff != "" { + t.Errorf("unexpected chunks (-want +got):\n%s", diff) + } + }) + + t.Run("Non-Live Mode (Aggregation)", func(t *testing.T) { + invCtx := icontext.NewInvocationContext(t.Context(), icontext.InvocationContextParams{ + InvocationID: "inv_1", + Agent: &mockAgent{name: "agent_1"}, + }) + + flow := &Flow{ + Tools: []tool.Tool{streamTool}, + } + + mergedEvent, err := flow.handleFunctionCalls(invCtx, toolsDict, resp, nil, nil) + if err != nil { + t.Fatal(err) + } + + if mergedEvent == nil { + t.Fatal("expected non-nil mergedEvent") + } + if len(mergedEvent.LLMResponse.Content.Parts) != 1 { + t.Fatalf("expected 1 part, got %d", len(mergedEvent.LLMResponse.Content.Parts)) + } + respPart := mergedEvent.LLMResponse.Content.Parts[0].FunctionResponse + if respPart == nil { + t.Fatal("expected FunctionResponse part") + } + result, ok := respPart.Response["result"].(string) + if !ok || result != "chunk 0chunk 1chunk 2" { + t.Errorf("unexpected result: %v", respPart.Response["result"]) + } + }) +} + +func TestHandleFunctionCalls_LiveControlPlane(t *testing.T) { + type Args struct { + DelayMS int `json:"delay_ms"` + } + + var cancelCount int + var cancelMu sync.Mutex + + handler := func(ctx agent.ToolContext, args Args) iter.Seq2[string, error] { + return func(yield func(string, error) bool) { + for i := 0; ; i++ { + time.Sleep(time.Millisecond) + select { + case <-ctx.Done(): + cancelMu.Lock() + cancelCount++ + cancelMu.Unlock() + return + default: + } + if !yield(fmt.Sprintf("number %d", i), nil) { + cancelMu.Lock() + cancelCount++ + cancelMu.Unlock() + return + } + } + } + } + + streamTool, err := functiontool.NewStreaming(functiontool.Config{ + Name: "count_forever", + Description: "counts indefinitely", + }, handler) + if err != nil { + t.Fatal(err) + } + + toolsDict := map[string]tool.Tool{ + "count_forever": streamTool, + } + + invCtx := icontext.NewInvocationContext(t.Context(), icontext.InvocationContextParams{ + InvocationID: "inv_1", + Agent: &mockAgent{name: "agent_1"}, + }) + + flow := &Flow{ + Tools: []tool.Tool{streamTool}, + } + + respStart := &model.LLMResponse{ + Content: &genai.Content{ + Parts: []*genai.Part{ + { + FunctionCall: &genai.FunctionCall{ + ID: "call_forever_1", + Name: "count_forever", + Args: map[string]any{"delay_ms": 1}, + }, + }, + { + FunctionCall: &genai.FunctionCall{ + ID: "call_forever_2", + Name: "count_forever", + Args: map[string]any{"delay_ms": 1}, + }, + }, + }, + Role: "model", + }, + } + + liveSess := newLiveSessionImpl() + liveSess.activeTools = make(map[string][]activeTask) + + go func() { + for range liveSess.inputCh { + } + }() + + _, err = flow.handleFunctionCalls(invCtx, toolsDict, respStart, nil, liveSess) + if err != nil { + t.Fatal(err) + } + + liveSess.mu.Lock() + tasks, exists := liveSess.activeTools["count_forever"] + liveSess.mu.Unlock() + if !exists || len(tasks) != 2 { + t.Fatalf("expected 2 active tasks, found: %v", tasks) + } + + respStop := &model.LLMResponse{ + Content: &genai.Content{ + Parts: []*genai.Part{ + { + FunctionCall: &genai.FunctionCall{ + ID: "call_stop_1", + Name: "stop_streaming", + Args: map[string]any{"function_name": "count_forever"}, + }, + }, + }, + Role: "model", + }, + } + + mergedStopEvent, err := flow.handleFunctionCalls(invCtx, toolsDict, respStop, nil, liveSess) + if err != nil { + t.Fatal(err) + } + + if mergedStopEvent == nil { + t.Fatal("expected non-nil mergedStopEvent") + } + respPart := mergedStopEvent.LLMResponse.Content.Parts[0].FunctionResponse + if respPart == nil { + t.Fatal("expected FunctionResponse part") + } + status, ok := respPart.Response["status"].(string) + if !ok || status != "Successfully stopped all running instances of count_forever" { + t.Errorf("unexpected stop status: %s", status) + } + + time.Sleep(50 * time.Millisecond) + + cancelMu.Lock() + gotCancels := cancelCount + cancelMu.Unlock() + if gotCancels != 2 { + t.Errorf("expected exactly 2 goroutines to be cancelled, got: %d", gotCancels) + } + + liveSess.mu.Lock() + tasksAfter := liveSess.activeTools["count_forever"] + liveSess.mu.Unlock() + if len(tasksAfter) != 0 { + t.Errorf("expected registry to be empty after cancellation, got: %v", tasksAfter) + } +} diff --git a/internal/plugininternal/plugin_manager.go b/internal/plugininternal/plugin_manager.go index 40b3403b1..c8ede432a 100644 --- a/internal/plugininternal/plugin_manager.go +++ b/internal/plugininternal/plugin_manager.go @@ -178,7 +178,7 @@ func (pm *PluginManager) RunAfterAgentCallback(cctx agent.CallbackContext) (*gen } // RunBeforeToolCallback runs the BeforeToolCallback for all plugins. -func (pm *PluginManager) RunBeforeToolCallback(ctx tool.Context, tool tool.Tool, args map[string]any) (map[string]any, error) { +func (pm *PluginManager) RunBeforeToolCallback(ctx agent.ToolContext, tool tool.Tool, args map[string]any) (map[string]any, error) { for _, plugin := range pm.plugins { callback := plugin.BeforeToolCallback() if callback != nil { @@ -195,7 +195,7 @@ func (pm *PluginManager) RunBeforeToolCallback(ctx tool.Context, tool tool.Tool, } // RunAfterToolCallback runs the AfterToolCallback for all plugins. -func (pm *PluginManager) RunAfterToolCallback(ctx tool.Context, tool tool.Tool, args, result map[string]any, err error) (map[string]any, error) { +func (pm *PluginManager) RunAfterToolCallback(ctx agent.ToolContext, tool tool.Tool, args, result map[string]any, err error) (map[string]any, error) { for _, plugin := range pm.plugins { callback := plugin.AfterToolCallback() if callback != nil { @@ -212,7 +212,7 @@ func (pm *PluginManager) RunAfterToolCallback(ctx tool.Context, tool tool.Tool, } // RunOnToolErrorCallback runs the OnToolErrorCallback for all plugins. -func (pm *PluginManager) RunOnToolErrorCallback(ctx tool.Context, tool tool.Tool, args map[string]any, err error) (map[string]any, error) { +func (pm *PluginManager) RunOnToolErrorCallback(ctx agent.ToolContext, tool tool.Tool, args map[string]any, err error) (map[string]any, error) { for _, plugin := range pm.plugins { callback := plugin.OnToolErrorCallback() if callback != nil { diff --git a/internal/telemetry/telemetry.go b/internal/telemetry/telemetry.go index 253083da8..3d159524a 100644 --- a/internal/telemetry/telemetry.go +++ b/internal/telemetry/telemetry.go @@ -42,12 +42,12 @@ const ( ) var ( - gcpVertexAgentToolCallArgsName = attribute.Key("gcp.vertex.agent.tool_call_args") - gcpVertexAgentEventID = attribute.Key("gcp.vertex.agent.event_id") - gcpVertexAgentToolResponseName = attribute.Key("gcp.vertex.agent.tool_response") - gcpVertexAgentInvocationID = attribute.Key("gcp.vertex.agent.invocation_id") - genAIUsageCacheReadInputTokens = attribute.Key("gen_ai.usage.cache_read.input_tokens") - genAIUsageExperimentalReasoningTokens = attribute.Key("gen_ai.usage.experimental.reasoning_tokens") + gcpVertexAgentToolCallArgsName = attribute.Key("gcp.vertex.agent.tool_call_args") + gcpVertexAgentEventID = attribute.Key("gcp.vertex.agent.event_id") + gcpVertexAgentToolResponseName = attribute.Key("gcp.vertex.agent.tool_response") + gcpVertexAgentInvocationID = attribute.Key("gcp.vertex.agent.invocation_id") + genAIUsageCacheReadInputTokens = attribute.Key("gen_ai.usage.cache_read.input_tokens") + genAIUsageReasoningOutputTokens = attribute.Key("gen_ai.usage.reasoning.output_tokens") // GCPVertexAgentOperationKey identifies the operation name for vendor-extension spans // where no semconv operation is appropriate (e.g., live_session). @@ -181,9 +181,12 @@ func TraceGenerateContentResult(span trace.Span, params TraceGenerateContentResu if params.Response.UsageMetadata != nil { span.SetAttributes( semconv.GenAIUsageInputTokens(int(params.Response.UsageMetadata.PromptTokenCount)), - semconv.GenAIUsageOutputTokens(int(params.Response.UsageMetadata.CandidatesTokenCount)), + // According to OpenTelemetry Semantic Conventions: + // https://github.com/open-telemetry/semantic-conventions/blob/v1.41.0/docs/registry/attributes/gen-ai.md + // gen_ai.usage.reasoning.output_tokens (ThoughtsTokenCount) SHOULD be included in gen_ai.usage.output_tokens. + semconv.GenAIUsageOutputTokens(int(params.Response.UsageMetadata.CandidatesTokenCount+params.Response.UsageMetadata.ThoughtsTokenCount)), genAIUsageCacheReadInputTokens.Int(int(params.Response.UsageMetadata.CachedContentTokenCount)), - genAIUsageExperimentalReasoningTokens.Int(int(params.Response.UsageMetadata.ThoughtsTokenCount)), + genAIUsageReasoningOutputTokens.Int(int(params.Response.UsageMetadata.ThoughtsTokenCount)), ) } } diff --git a/internal/telemetry/telemetry_test.go b/internal/telemetry/telemetry_test.go index d775a4f3c..8b0c17dfe 100644 --- a/internal/telemetry/telemetry_test.go +++ b/internal/telemetry/telemetry_test.go @@ -129,7 +129,7 @@ func TestInvokeAgent(t *testing.T) { { name: "Success", resultParams: TraceAgentResultParams{ - ResponseEvent: session.NewEvent("test-invocation-id"), + ResponseEvent: session.NewEventWithContext(context.Background(), "test-invocation-id"), }, wantName: "invoke_agent test-agent", wantStatus: codes.Unset, @@ -230,9 +230,9 @@ func TestGenerateContent(t *testing.T) { semconv.GenAIOperationNameKey: "generate_content", semconv.GenAIRequestModelKey: "test-model", semconv.GenAIUsageInputTokensKey: "10", - semconv.GenAIUsageOutputTokensKey: "20", + semconv.GenAIUsageOutputTokensKey: "35", genAIUsageCacheReadInputTokens: "5", - genAIUsageExperimentalReasoningTokens: "15", + genAIUsageReasoningOutputTokens: "15", semconv.GenAIResponseFinishReasonsKey: "[\"STOP\"]", gcpVertexAgentInvocationID: invocationID, }, diff --git a/internal/toolinternal/context.go b/internal/toolinternal/context.go deleted file mode 100644 index 541b061d0..000000000 --- a/internal/toolinternal/context.go +++ /dev/null @@ -1,136 +0,0 @@ -// Copyright 2025 Google LLC -// -// Licensed under the Apache License, Version 2.0 (the "License"); -// you may not use this file except in compliance with the License. -// You may obtain a copy of the License at -// -// http://www.apache.org/licenses/LICENSE-2.0 -// -// Unless required by applicable law or agreed to in writing, software -// distributed under the License is distributed on an "AS IS" BASIS, -// WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. -// See the License for the specific language governing permissions and -// limitations under the License. - -package toolinternal - -import ( - "context" - "fmt" - - "github.com/google/uuid" - "google.golang.org/genai" - - "google.golang.org/adk/agent" - "google.golang.org/adk/artifact" - contextinternal "google.golang.org/adk/internal/context" - "google.golang.org/adk/memory" - "google.golang.org/adk/session" - "google.golang.org/adk/tool" - "google.golang.org/adk/tool/toolconfirmation" -) - -type internalArtifacts struct { - agent.Artifacts - eventActions *session.EventActions -} - -func (ia *internalArtifacts) Save(ctx context.Context, name string, data *genai.Part) (*artifact.SaveResponse, error) { - resp, err := ia.Artifacts.Save(ctx, name, data) - if err != nil { - return resp, err - } - if ia.eventActions != nil { - if ia.eventActions.ArtifactDelta == nil { - ia.eventActions.ArtifactDelta = make(map[string]int64) - } - // TODO: RWLock, check the version stored is newer in case multiple tools save the same file. - ia.eventActions.ArtifactDelta[name] = resp.Version - } - return resp, nil -} - -func NewToolContext(ctx agent.InvocationContext, functionCallID string, actions *session.EventActions, confirmation *toolconfirmation.ToolConfirmation) tool.Context { - if functionCallID == "" { - functionCallID = uuid.NewString() - } - if actions == nil { - actions = &session.EventActions{StateDelta: make(map[string]any)} - } - if actions.StateDelta == nil { - actions.StateDelta = make(map[string]any) - } - if actions.ArtifactDelta == nil { - actions.ArtifactDelta = make(map[string]int64) - } - cbCtx := contextinternal.NewCallbackContextWithDelta(ctx, actions.StateDelta, actions.ArtifactDelta) - - return &toolContext{ - CallbackContext: cbCtx, - invocationContext: ctx, - functionCallID: functionCallID, - eventActions: actions, - artifacts: &internalArtifacts{ - Artifacts: ctx.Artifacts(), - eventActions: actions, - }, - toolConfirmation: confirmation, - } -} - -type toolContext struct { - agent.CallbackContext - invocationContext agent.InvocationContext - functionCallID string - eventActions *session.EventActions - artifacts *internalArtifacts - toolConfirmation *toolconfirmation.ToolConfirmation -} - -func (c *toolContext) Artifacts() agent.Artifacts { - return c.artifacts -} - -func (c *toolContext) FunctionCallID() string { - return c.functionCallID -} - -func (c *toolContext) Actions() *session.EventActions { - return c.eventActions -} - -func (c *toolContext) AgentName() string { - return c.invocationContext.Agent().Name() -} - -func (c *toolContext) SearchMemory(ctx context.Context, query string) (*memory.SearchResponse, error) { - if c.invocationContext.Memory() == nil { - return nil, fmt.Errorf("memory service is not set") - } - return c.invocationContext.Memory().SearchMemory(ctx, query) -} - -func (c *toolContext) ToolConfirmation() *toolconfirmation.ToolConfirmation { - return c.toolConfirmation -} - -func (c *toolContext) RequestConfirmation(hint string, payload any) error { - if c.functionCallID == "" { - return fmt.Errorf("error function call id not set when requesting confirmation for tool") - } - if c.eventActions.RequestedToolConfirmations == nil { - c.eventActions.RequestedToolConfirmations = make(map[string]toolconfirmation.ToolConfirmation) - } - c.eventActions.RequestedToolConfirmations[c.functionCallID] = toolconfirmation.ToolConfirmation{ - Hint: hint, - Confirmed: false, - Payload: payload, - } - // SkipSummarization stops the agent loop after this tool call. Without it, - // the function response event becomes lastEvent and IsFinalResponse() returns - // false (hasFunctionResponses == true), causing the loop to continue. - // This matches the behavior of the built-in RequireConfirmation path in - // functiontool (function.go). - c.eventActions.SkipSummarization = true - return nil -} diff --git a/internal/toolinternal/tool.go b/internal/toolinternal/tool.go index 40aa63fd2..1655f0844 100644 --- a/internal/toolinternal/tool.go +++ b/internal/toolinternal/tool.go @@ -16,8 +16,11 @@ package toolinternal import ( + "iter" + "google.golang.org/genai" + "google.golang.org/adk/agent" "google.golang.org/adk/model" "google.golang.org/adk/tool" ) @@ -25,9 +28,15 @@ import ( type FunctionTool interface { tool.Tool Declaration() *genai.FunctionDeclaration - Run(ctx tool.Context, args any) (result map[string]any, err error) + Run(ctx agent.ToolContext, args any) (result map[string]any, err error) +} + +type StreamingFunctionTool interface { + tool.Tool + Declaration() *genai.FunctionDeclaration + RunStream(ctx agent.ToolContext, args any) iter.Seq2[string, error] } type RequestProcessor interface { - ProcessRequest(ctx tool.Context, req *model.LLMRequest) error + ProcessRequest(ctx agent.ToolContext, req *model.LLMRequest) error } diff --git a/internal/utils/utils.go b/internal/utils/utils.go index 3d29baa08..895229f05 100644 --- a/internal/utils/utils.go +++ b/internal/utils/utils.go @@ -15,13 +15,14 @@ package utils import ( + "context" "strings" - "github.com/google/uuid" "google.golang.org/genai" "google.golang.org/adk/agent" "google.golang.org/adk/model" + "google.golang.org/adk/platform" "google.golang.org/adk/session" ) @@ -33,17 +34,19 @@ const afFunctionCallIDPrefix = "adk-" // Since the ID field is optional, some models don't fill the field, but // the LLMAgent depends on the IDs to map FunctionCall and FunctionResponse events // in the event stream. -func PopulateClientFunctionCallID(c *genai.Content) { +func PopulateClientFunctionCallID(ctx context.Context, c *genai.Content) { for _, fn := range FunctionCalls(c) { if fn.ID == "" { - fn.ID = GenerateFunctionCallID() + fn.ID = GenerateFunctionCallID(ctx) } } } -// GenerateFunctionCallID generates a new function call ID. -func GenerateFunctionCallID() string { - return afFunctionCallIDPrefix + uuid.NewString() +// GenerateFunctionCallID generates a new function call ID. The ID is obtained +// through the platform package, so a UUID provider installed on ctx (see +// platform.WithUUIDProvider) controls it. +func GenerateFunctionCallID(ctx context.Context) string { + return afFunctionCallIDPrefix + platform.NewUUID(ctx) } // RemoveClientFunctionCallID removes the function call ID field that was set diff --git a/internal/utils/utils_test.go b/internal/utils/utils_test.go index 9b5144d23..3594b132c 100644 --- a/internal/utils/utils_test.go +++ b/internal/utils/utils_test.go @@ -14,8 +14,57 @@ package utils_test -import "testing" +import ( + "context" + "strings" + "testing" -func TestNothing(t *testing.T) { - // To make it buildable. + "google.golang.org/genai" + + "google.golang.org/adk/internal/utils" + "google.golang.org/adk/platform" +) + +func TestGenerateFunctionCallIDUsesProvider(t *testing.T) { + ctx := platform.WithUUIDProvider(context.Background(), func() string { return "fixed" }) + + got := utils.GenerateFunctionCallID(ctx) + + // The generated ID must carry the "adk-" prefix that RemoveClientFunctionCallID + // relies on, and must incorporate the value from the installed provider. + if !strings.HasPrefix(got, "adk-") { + t.Errorf("GenerateFunctionCallID() = %q, want \"adk-\" prefix", got) + } + if !strings.HasSuffix(got, "fixed") { + t.Errorf("GenerateFunctionCallID() = %q, want it to use the provider value %q", got, "fixed") + } +} + +func TestGenerateFunctionCallIDDefaultIsUnique(t *testing.T) { + first := utils.GenerateFunctionCallID(context.Background()) + second := utils.GenerateFunctionCallID(context.Background()) + + if first == second { + t.Errorf("GenerateFunctionCallID() returned %q twice; want unique values", first) + } +} + +func TestPopulateClientFunctionCallIDUsesProvider(t *testing.T) { + ctx := platform.WithUUIDProvider(context.Background(), func() string { return "generated" }) + + content := &genai.Content{ + Parts: []*genai.Part{ + {FunctionCall: &genai.FunctionCall{Name: "needs_id"}}, + {FunctionCall: &genai.FunctionCall{ID: "keep", Name: "has_id"}}, + }, + } + + utils.PopulateClientFunctionCallID(ctx, content) + + if got := content.Parts[0].FunctionCall.ID; got != "adk-generated" { + t.Errorf("empty function call ID = %q, want %q", got, "adk-generated") + } + if got := content.Parts[1].FunctionCall.ID; got != "keep" { + t.Errorf("preset function call ID = %q, want it left untouched (%q)", got, "keep") + } } diff --git a/internal/version/version.go b/internal/version/version.go index 2ff67683e..bee420370 100644 --- a/internal/version/version.go +++ b/internal/version/version.go @@ -15,4 +15,4 @@ package version // Version exposes the current ADK Go version, used for llm request tagging -const Version = "1.0.0" +const Version = "1.2.0" diff --git a/memory/vertexai/vertexai.go b/memory/vertexai/vertexai.go new file mode 100644 index 000000000..5355fc23e --- /dev/null +++ b/memory/vertexai/vertexai.go @@ -0,0 +1,95 @@ +// Copyright 2026 Google LLC +// +// Licensed under the Apache License, Version 2.0 (the "License"); +// you may not use this file except in compliance with the License. +// You may obtain a copy of the License at +// +// http://www.apache.org/licenses/LICENSE-2.0 +// +// Unless required by applicable law or agreed to in writing, software +// distributed under the License is distributed on an "AS IS" BASIS, +// WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. +// See the License for the specific language governing permissions and +// limitations under the License. + +// Package vertexai provides support for using MemoryBank provided by VertexAI +package vertexai + +import ( + "context" + "fmt" + "time" + + "google.golang.org/adk/memory" + "google.golang.org/adk/session" + vertexaiutil "google.golang.org/adk/util/vertexai" +) + +type vertexAIService struct { + client *vertexAIClient + stateKeySessionLastUpdateTime string +} + +// ServiceConfig allows you to specify the instance of MemoryBank (by specifying an AgentEngine instance). +// Specifies also a way to convert session events to memories (see StateKeySessionLastUpdateTime) +type ServiceConfig struct { + vertexaiutil.AgentEngineData + // StateKeySessionLastUpdateTime controlls the process of the generation of memories. + // If set to "", the whole session is used to generate the memories. + // If provided, the value is treated as a key for Session State. Retrieved value (time.Time is expected) is used to filter the events to the most recent ones. + // Warning! The value for State under this key should be set as soon as possible (for instance, during the BeforeRunCallback) + StateKeySessionLastUpdateTime string + WaitForCompletion bool +} + +// NewService creates a new instance of Memory Service supported by VertexAI MemoryBank. +func NewService(ctx context.Context, config *ServiceConfig) (memory.Service, error) { + client, err := newVertexAIClient(ctx, &vertexAIClientConfig{ + AgentEngineData: config.AgentEngineData, + waitForCompletion: config.WaitForCompletion, + }) + if err != nil { + return nil, fmt.Errorf("newVertexAIClient failed: %w", err) + } + return &vertexAIService{ + client: client, + stateKeySessionLastUpdateTime: config.StateKeySessionLastUpdateTime, + }, nil +} + +var _ memory.Service = &vertexAIService{} + +// AddSessionToMemory implements [memory.Service]. +func (v *vertexAIService) AddSessionToMemory(ctx context.Context, s session.Session) error { + var err error + if v.stateKeySessionLastUpdateTime == "" { + // add the whole session + err = v.client.addWholeSession(ctx, s) + if err != nil { + return fmt.Errorf("v.client.addWholeSession failed: %w", err) + } + return nil + } + + // add only events newer than given + t, err := s.State().Get(v.stateKeySessionLastUpdateTime) + if err != nil { + return fmt.Errorf("state.Get(%s) failed: %w", v.stateKeySessionLastUpdateTime, err) + } + + tm, ok := t.(time.Time) + if !ok { + return fmt.Errorf("want type time.Time, got : %T (%v)", t, t) + } + err = v.client.addEventsNewerThan(ctx, s, tm) + if err != nil { + return fmt.Errorf("v.client.addEventsNewerThan failed: %w", err) + } + + return err +} + +// SearchMemory implements [memory.Service]. +func (v *vertexAIService) SearchMemory(ctx context.Context, req *memory.SearchRequest) (*memory.SearchResponse, error) { + return v.client.searchMemory(ctx, req) +} diff --git a/memory/vertexai/vertexai_client.go b/memory/vertexai/vertexai_client.go new file mode 100644 index 000000000..4926337fb --- /dev/null +++ b/memory/vertexai/vertexai_client.go @@ -0,0 +1,139 @@ +// Copyright 2026 Google LLC +// +// Licensed under the Apache License, Version 2.0 (the "License"); +// you may not use this file except in compliance with the License. +// You may obtain a copy of the License at +// +// http://www.apache.org/licenses/LICENSE-2.0 +// +// Unless required by applicable law or agreed to in writing, software +// distributed under the License is distributed on an "AS IS" BASIS, +// WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. +// See the License for the specific language governing permissions and +// limitations under the License. + +package vertexai + +import ( + "context" + "fmt" + "time" + + "google.golang.org/api/option" + "google.golang.org/genai" + "google.golang.org/protobuf/types/known/timestamppb" + + "google.golang.org/adk/memory" + "google.golang.org/adk/session" + aiplatformutil "google.golang.org/adk/util/aiplatform" + vertexaiutil "google.golang.org/adk/util/vertexai" + + aiplatform "cloud.google.com/go/aiplatform/apiv1beta1" + "cloud.google.com/go/aiplatform/apiv1beta1/aiplatformpb" +) + +type vertexAIClient struct { + config vertexAIClientConfig + client *aiplatform.MemoryBankClient + agentEngineData *vertexaiutil.AgentEngineData + parent string +} + +type vertexAIClientConfig struct { + vertexaiutil.AgentEngineData + waitForCompletion bool +} + +func newVertexAIClient(ctx context.Context, config *vertexAIClientConfig) (*vertexAIClient, error) { + if config == nil { + return nil, fmt.Errorf("config cannot be nil") + } + c, err := aiplatform.NewMemoryBankClient(ctx, option.WithEndpoint(aiplatformutil.HostPortURL(config.Location))) + if err != nil { + return nil, fmt.Errorf("aiplatform.NewMemoryBankClient failed: %w", err) + } + return &vertexAIClient{ + config: *config, + client: c, + agentEngineData: &config.AgentEngineData, + parent: vertexaiutil.AgentEngineResource(&config.AgentEngineData), + }, nil +} + +// addWholeSession adds the whole session to the Memory +func (v *vertexAIClient) addWholeSession(ctx context.Context, s session.Session) error { + return v.addSession(ctx, s, nil) +} + +// addEventsNewerThan uses time to filter out the old event. The new ones are used to generate memories +func (v *vertexAIClient) addEventsNewerThan(ctx context.Context, s session.Session, start time.Time) error { + return v.addSession(ctx, s, &start) +} + +// addSession adds the whole session or just events created after `start` to the Memory +func (v *vertexAIClient) addSession(ctx context.Context, s session.Session, start *time.Time) error { + sr := vertexaiutil.SessionResource(v.agentEngineData, s.ID()) + + vss := &aiplatformpb.GenerateMemoriesRequest_VertexSessionSource{ + Session: sr, + } + if start != nil { + vss.StartTime = timestamppb.New(*start) + } + req := &aiplatformpb.GenerateMemoriesRequest{ + Parent: v.parent, + Source: &aiplatformpb.GenerateMemoriesRequest_VertexSessionSource_{VertexSessionSource: vss}, + Scope: createUserScope(s.UserID()), + } + + op, err := v.client.GenerateMemories(ctx, req) + if err != nil { + return fmt.Errorf("v.client.GenerateMemories failed: %w", err) + } + if v.config.waitForCompletion { + _, err = op.Wait(ctx) + if err != nil && err.Error() == "unsupported result type : " { + // accept it + err = nil + } + if err != nil { + return fmt.Errorf("op.Wait for GenerateMemories (whole session) failed: %w", err) + } + } + return nil +} + +// searchMemory uses provided query to find the relevant memories for the given user +func (v *vertexAIClient) searchMemory(ctx context.Context, req *memory.SearchRequest) (*memory.SearchResponse, error) { + r, err := v.client.RetrieveMemories(ctx, + &aiplatformpb.RetrieveMemoriesRequest{ + RetrievalParams: &aiplatformpb.RetrieveMemoriesRequest_SimilaritySearchParams_{ + SimilaritySearchParams: &aiplatformpb.RetrieveMemoriesRequest_SimilaritySearchParams{ + SearchQuery: req.Query, + }, + }, + Parent: v.parent, + Scope: createUserScope(req.UserID), + }, + ) + if err != nil { + return nil, fmt.Errorf("v.client.RetrieveMemories failed: %w", err) + } + + res := &memory.SearchResponse{ + Memories: []memory.Entry{}, + } + + for _, m := range r.RetrievedMemories { + res.Memories = append(res.Memories, memory.Entry{ + Content: genai.NewContentFromText(m.Memory.Fact, genai.RoleUser), + }) + } + + return res, nil +} + +// Scope is used to structure the information in MemoryBank. Here we use only the user scope +func createUserScope(userID string) map[string]string { + return map[string]string{"user_id": userID} +} diff --git a/model/gemini/gemini.go b/model/gemini/gemini.go index cb1e79ec0..39ce3ba8d 100644 --- a/model/gemini/gemini.go +++ b/model/gemini/gemini.go @@ -196,4 +196,8 @@ func (m *geminiModel) GetGoogleLLMVariant() genai.Backend { return m.client.ClientConfig().Backend } +func (m *geminiModel) Client() *genai.Client { + return m.client +} + var _ googlellm.GoogleLLM = &geminiModel{} diff --git a/model/llm.go b/model/llm.go index 5b706eb2d..8cb2d4b29 100644 --- a/model/llm.go +++ b/model/llm.go @@ -42,13 +42,15 @@ type LLMRequest struct { // LLMResponse is the raw LLM response. // It provides the first candidate response from the model if available. type LLMResponse struct { - Content *genai.Content - CitationMetadata *genai.CitationMetadata - GroundingMetadata *genai.GroundingMetadata - UsageMetadata *genai.GenerateContentResponseUsageMetadata - CustomMetadata map[string]any - LogprobsResult *genai.LogprobsResult - ModelVersion string + Content *genai.Content + CitationMetadata *genai.CitationMetadata + GroundingMetadata *genai.GroundingMetadata + UsageMetadata *genai.GenerateContentResponseUsageMetadata + CustomMetadata map[string]any + LogprobsResult *genai.LogprobsResult + InputTranscription *genai.Transcription + OutputTranscription *genai.Transcription + ModelVersion string // Partial indicates whether the content is part of a unfinished content stream. // Only used for streaming mode and when the content is plain text. // The Runner fully processes only the final non-partial event, partial @@ -61,13 +63,16 @@ type LLMResponse struct { // Usually it is due to user interruption during a bidi streaming. Interrupted bool - // Live-only: transcription of user audio input / model audio output. - // Populated by the model connector (e.g. gemini_live.go) from the Live API's - // ServerContent.InputTranscription / OutputTranscription fields. - InputTranscription *genai.Transcription - OutputTranscription *genai.Transcription + // SessionResumptionHandle is the upstream (v1.5.0) live resumption handle. + SessionResumptionHandle string - // Live-only: session resumption state update from the server. + // SessionResumptionUpdate is the fork's structured session-resumption + // carrier (Live-only). Two carriers coexist by design — do not remove + // either: the upstream live engine (llminternal Flow.RunLive, fed by + // the googlellm connection) reads the plain SessionResumptionHandle + // string above, while the hulilabs liveflow engine + // (internal/llminternal/liveflow, fed by model/gemini) consumes this + // full update so it can also honor Resumable=false handle invalidation. SessionResumptionUpdate *genai.LiveServerSessionResumptionUpdate // Live-only: GoAway signal indicating the server wants the client to reconnect. GoAway *genai.LiveServerGoAway diff --git a/platform/doc.go b/platform/doc.go new file mode 100644 index 000000000..7c32a17ff --- /dev/null +++ b/platform/doc.go @@ -0,0 +1,29 @@ +// Copyright 2026 Google LLC +// +// Licensed under the Apache License, Version 2.0 (the "License"); +// you may not use this file except in compliance with the License. +// You may obtain a copy of the License at +// +// http://www.apache.org/licenses/LICENSE-2.0 +// +// Unless required by applicable law or agreed to in writing, software +// distributed under the License is distributed on an "AS IS" BASIS, +// WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. +// See the License for the specific language governing permissions and +// limitations under the License. + +// Package platform provides seams for overriding system operations, such as +// reading the current time and generating unique IDs. +// +// By default these operations use the wall clock (time.Now) and random UUIDs. +// Users that need custom behavior can install their own providers on a +// context.Context with WithTimeProvider and WithUUIDProvider. ADK reads the +// current time and new IDs through Now and NewUUID, so an installed provider +// transparently controls the timestamps and identifiers that end up in events +// and sessions. +// +// Providers are carried explicitly on a context.Context. Carrying the +// provider on the context (rather than in a package-level variable) also +// keeps it isolated to a single call tree, which is what makes concurrent +// runs with independent providers safe. +package platform diff --git a/platform/time.go b/platform/time.go new file mode 100644 index 000000000..c5b9c0033 --- /dev/null +++ b/platform/time.go @@ -0,0 +1,50 @@ +// Copyright 2026 Google LLC +// +// Licensed under the Apache License, Version 2.0 (the "License"); +// you may not use this file except in compliance with the License. +// You may obtain a copy of the License at +// +// http://www.apache.org/licenses/LICENSE-2.0 +// +// Unless required by applicable law or agreed to in writing, software +// distributed under the License is distributed on an "AS IS" BASIS, +// WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. +// See the License for the specific language governing permissions and +// limitations under the License. + +package platform + +import ( + "context" + "time" +) + +// TimeProvider returns the current time. The default behavior is the wall +// clock (time.Now); callers can install a custom provider on a context +// with WithTimeProvider. +type TimeProvider func() time.Time + +// timeProviderKey is the context key under which a TimeProvider is stored. +type timeProviderKey struct{} + +// WithTimeProvider returns a copy of ctx that carries provider. Calls to Now +// with the returned context, or any context derived from it, use provider +// instead of the wall clock. A nil provider is ignored by Now, which then +// falls back to time.Now. +func WithTimeProvider(ctx context.Context, provider TimeProvider) context.Context { + return context.WithValue(ctx, timeProviderKey{}, provider) +} + +// Now returns the current time. If ctx carries a TimeProvider installed with +// WithTimeProvider, that provider is used; otherwise Now falls back to +// time.Now. +// +// Now is the analog of google.adk.platform.time.get_time in adk-python. +func Now(ctx context.Context) time.Time { + if ctx != nil { + if p, ok := ctx.Value(timeProviderKey{}).(TimeProvider); ok && p != nil { + return p() + } + } + return time.Now() +} diff --git a/platform/time_test.go b/platform/time_test.go new file mode 100644 index 000000000..65381e099 --- /dev/null +++ b/platform/time_test.go @@ -0,0 +1,109 @@ +// Copyright 2026 Google LLC +// +// Licensed under the Apache License, Version 2.0 (the "License"); +// you may not use this file except in compliance with the License. +// You may obtain a copy of the License at +// +// http://www.apache.org/licenses/LICENSE-2.0 +// +// Unless required by applicable law or agreed to in writing, software +// distributed under the License is distributed on an "AS IS" BASIS, +// WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. +// See the License for the specific language governing permissions and +// limitations under the License. + +package platform_test + +import ( + "context" + "testing" + "time" + + "google.golang.org/adk/platform" +) + +func TestNowDefaultUsesWallClock(t *testing.T) { + before := time.Now() + got := platform.Now(context.Background()) + after := time.Now() + + if got.Before(before) || got.After(after) { + t.Errorf("Now() = %v, want within [%v, %v]", got, before, after) + } +} + +func TestNowNilContext(t *testing.T) { + // Now must tolerate a nil context and fall back to the wall clock. + var nilCtx context.Context + before := time.Now() + got := platform.Now(nilCtx) + after := time.Now() + + if got.Before(before) || got.After(after) { + t.Errorf("Now(nil) = %v, want within [%v, %v]", got, before, after) + } +} + +func TestWithTimeProviderOverridesNow(t *testing.T) { + fixed := time.Date(2024, time.January, 2, 3, 4, 5, 0, time.UTC) + ctx := platform.WithTimeProvider(context.Background(), func() time.Time { return fixed }) + + if got := platform.Now(ctx); !got.Equal(fixed) { + t.Errorf("Now() = %v, want %v", got, fixed) + } +} + +func TestWithTimeProviderDerivedContext(t *testing.T) { + fixed := time.Date(2024, time.January, 2, 3, 4, 5, 0, time.UTC) + ctx := platform.WithTimeProvider(context.Background(), func() time.Time { return fixed }) + + // A context derived from the provider context must keep the provider. + derived, cancel := context.WithCancel(ctx) + defer cancel() + + if got := platform.Now(derived); !got.Equal(fixed) { + t.Errorf("Now() on derived context = %v, want %v", got, fixed) + } +} + +func TestWithTimeProviderNilFallsBack(t *testing.T) { + ctx := platform.WithTimeProvider(context.Background(), nil) + + before := time.Now() + got := platform.Now(ctx) + after := time.Now() + + if got.Before(before) || got.After(after) { + t.Errorf("Now() with nil provider = %v, want within [%v, %v]", got, before, after) + } +} + +func TestWithTimeProviderNestedOverride(t *testing.T) { + outer := time.Date(2024, time.January, 1, 0, 0, 0, 0, time.UTC) + inner := time.Date(2025, time.June, 7, 8, 9, 10, 0, time.UTC) + + ctx := platform.WithTimeProvider(context.Background(), func() time.Time { return outer }) + ctx = platform.WithTimeProvider(ctx, func() time.Time { return inner }) + + if got := platform.Now(ctx); !got.Equal(inner) { + t.Errorf("Now() = %v, want innermost provider value %v", got, inner) + } +} + +func TestWithTimeProviderIsCalledEachTime(t *testing.T) { + var calls int + ctx := platform.WithTimeProvider(context.Background(), func() time.Time { + calls++ + return time.Unix(int64(calls), 0).UTC() + }) + + first := platform.Now(ctx) + second := platform.Now(ctx) + + if first.Equal(second) { + t.Errorf("Now() returned the same time %v on consecutive calls; provider should be invoked each call", first) + } + if calls != 2 { + t.Errorf("provider called %d times, want 2", calls) + } +} diff --git a/platform/uuid.go b/platform/uuid.go new file mode 100644 index 000000000..7de33868c --- /dev/null +++ b/platform/uuid.go @@ -0,0 +1,52 @@ +// Copyright 2026 Google LLC +// +// Licensed under the Apache License, Version 2.0 (the "License"); +// you may not use this file except in compliance with the License. +// You may obtain a copy of the License at +// +// http://www.apache.org/licenses/LICENSE-2.0 +// +// Unless required by applicable law or agreed to in writing, software +// distributed under the License is distributed on an "AS IS" BASIS, +// WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. +// See the License for the specific language governing permissions and +// limitations under the License. + +package platform + +import ( + "context" + + "github.com/google/uuid" +) + +// UUIDProvider returns a new unique identifier. The default behavior is to +// return a random UUIDv4 (uuid.NewString); callers can install a custom +// provider on a context with WithUUIDProvider. +type UUIDProvider func() string + +// uuidProviderKey is the context key under which a UUIDProvider is stored. +type uuidProviderKey struct{} + +// WithUUIDProvider returns a copy of ctx that carries provider. Calls to +// NewUUID with the returned context, or any context derived from it, use +// provider instead of generating a random UUID. A nil provider is ignored by +// NewUUID, which then falls back to uuid.NewString. +func WithUUIDProvider(ctx context.Context, provider UUIDProvider) context.Context { + return context.WithValue(ctx, uuidProviderKey{}, provider) +} + +// NewUUID returns a new unique identifier. If ctx carries a UUIDProvider +// installed with WithUUIDProvider, that provider is used; otherwise NewUUID +// falls back to a random UUIDv4 (uuid.NewString). +// +// NewUUID is the analog of google.adk.platform.uuid.new_uuid in adk-python, +// with the provider read from ctx rather than from a contextvar. +func NewUUID(ctx context.Context) string { + if ctx != nil { + if p, ok := ctx.Value(uuidProviderKey{}).(UUIDProvider); ok && p != nil { + return p() + } + } + return uuid.NewString() +} diff --git a/platform/uuid_test.go b/platform/uuid_test.go new file mode 100644 index 000000000..cadec2b3a --- /dev/null +++ b/platform/uuid_test.go @@ -0,0 +1,89 @@ +// Copyright 2026 Google LLC +// +// Licensed under the Apache License, Version 2.0 (the "License"); +// you may not use this file except in compliance with the License. +// You may obtain a copy of the License at +// +// http://www.apache.org/licenses/LICENSE-2.0 +// +// Unless required by applicable law or agreed to in writing, software +// distributed under the License is distributed on an "AS IS" BASIS, +// WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. +// See the License for the specific language governing permissions and +// limitations under the License. + +package platform_test + +import ( + "context" + "testing" + + "github.com/google/uuid" + + "google.golang.org/adk/platform" +) + +func TestNewUUIDDefaultIsRandomAndValid(t *testing.T) { + first := platform.NewUUID(context.Background()) + second := platform.NewUUID(context.Background()) + + if _, err := uuid.Parse(first); err != nil { + t.Errorf("NewUUID() = %q, not a valid UUID: %v", first, err) + } + if first == second { + t.Errorf("NewUUID() returned the same value %q twice; want random values", first) + } +} + +func TestNewUUIDNilContext(t *testing.T) { + // NewUUID must tolerate a nil context and fall back to a random UUID. + var nilCtx context.Context + got := platform.NewUUID(nilCtx) + if _, err := uuid.Parse(got); err != nil { + t.Errorf("NewUUID(nil) = %q, not a valid UUID: %v", got, err) + } +} + +func TestWithUUIDProviderOverridesNewUUID(t *testing.T) { + var n int + ctx := platform.WithUUIDProvider(context.Background(), func() string { + n++ + return "id-" + string(rune('0'+n)) + }) + + want := []string{"id-1", "id-2", "id-3"} + for i, w := range want { + if got := platform.NewUUID(ctx); got != w { + t.Errorf("NewUUID() call %d = %q, want %q", i+1, got, w) + } + } +} + +func TestWithUUIDProviderDerivedContext(t *testing.T) { + ctx := platform.WithUUIDProvider(context.Background(), func() string { return "fixed" }) + + derived, cancel := context.WithCancel(ctx) + defer cancel() + + if got := platform.NewUUID(derived); got != "fixed" { + t.Errorf("NewUUID() on derived context = %q, want %q", got, "fixed") + } +} + +func TestWithUUIDProviderNilFallsBack(t *testing.T) { + ctx := platform.WithUUIDProvider(context.Background(), nil) + + got := platform.NewUUID(ctx) + if _, err := uuid.Parse(got); err != nil { + t.Errorf("NewUUID() with nil provider = %q, not a valid UUID: %v", got, err) + } +} + +func TestWithUUIDProviderNestedOverride(t *testing.T) { + ctx := platform.WithUUIDProvider(context.Background(), func() string { return "outer" }) + ctx = platform.WithUUIDProvider(ctx, func() string { return "inner" }) + + if got := platform.NewUUID(ctx); got != "inner" { + t.Errorf("NewUUID() = %q, want innermost provider value %q", got, "inner") + } +} diff --git a/plugin/functioncallmodifier/plugin_test.go b/plugin/functioncallmodifier/plugin_test.go index 4b16dd998..3eca6955f 100644 --- a/plugin/functioncallmodifier/plugin_test.go +++ b/plugin/functioncallmodifier/plugin_test.go @@ -41,7 +41,7 @@ type SimpleArgs struct { Num int } -func okFunc(_ tool.Context, _ SimpleArgs) (string, error) { +func okFunc(_ agent.ToolContext, _ SimpleArgs) (string, error) { return "ok", nil } diff --git a/plugin/functioncallmodifier/testdata/TestPluginCallbackIntegration_transfer_to_agent_tool.httprr b/plugin/functioncallmodifier/testdata/TestPluginCallbackIntegration_transfer_to_agent_tool.httprr index c55277002..eac8c191d 100644 --- a/plugin/functioncallmodifier/testdata/TestPluginCallbackIntegration_transfer_to_agent_tool.httprr +++ b/plugin/functioncallmodifier/testdata/TestPluginCallbackIntegration_transfer_to_agent_tool.httprr @@ -1,12 +1,12 @@ httprr trace v1 -1870 1661 +1892 1661 POST https://generativelanguage.googleapis.com/v1beta/models/gemini-2.5-flash:generateContent HTTP/1.1 Host: generativelanguage.googleapis.com User-Agent: Go-http-client/1.1 -Content-Length: 1637 +Content-Length: 1659 Content-Type: application/json -{"contents":[{"parts":[{"text":"Can you add 2 and 2?"}],"role":"user"}],"generationConfig":{},"systemInstruction":{"parts":[{"text":"You are a transfer agent. You can transfer to other agents using your tools.\n\nYou are an agent. Your internal name is \"transfer_agent\". The description about you is \"transfer agent\".\n\n\nYou have a list of other agents to transfer to:\n\n\nAgent name: calculator\nAgent description: calculator agent\n Skills: add, subtract, multiply, divide\n\n\nIf you are the best to answer the question according to your description,\nyou can answer it.\n\nIf another agent is better for answering the question according to its\ndescription, call `transfer_to_agent` function to transfer the question to that\nagent. When transferring, do not generate any text other than the function\ncall.\n\n**NOTE**: the only available agents for `transfer_to_agent` function are\n`calculator`.\n"}],"role":"user"},"tools":[{"functionDeclarations":[{"description":"This tool can now optionally accept skill_id and rationale parameters to guide skill-based orchestration. Transfer the question to another agent.\nThis tool hands off control to another agent when it's more suitable to answer the user's question according to the agent's description.","name":"transfer_to_agent","parameters":{"properties":{"agent_name":{"description":"the agent name to transfer to","type":"string"},"rationale":{"description":"The reasoning behind selecting this agent and skill.","type":"STRING"},"skill_id":{"description":"The specific skill to be utilized by the agent.","type":"STRING"}},"required":["agent_name"],"type":"object"}}]}]}HTTP/2.0 200 OK +{"contents":[{"parts":[{"text":"Can you add 2 and 2?"}],"role":"user"}],"generationConfig":{},"systemInstruction":{"parts":[{"text":"You are a transfer agent. You can transfer to other agents using your tools.\n\nYou are an agent. Your internal name is \"transfer_agent\". The description about you is \"transfer agent\".\n\n\nYou have a list of other agents to transfer to:\n\n\nAgent name: calculator\nAgent description: calculator agent\n Skills: add, subtract, multiply, divide\n\n\nIf you are the best to answer the question according to your description,\nyou can answer it.\n\nIf another agent is better for answering the question according to its\ndescription, call `transfer_to_agent` function to transfer the question to that\nagent. When transferring, do not generate any text other than the function\ncall.\n\n**NOTE**: the only available agents for `transfer_to_agent` function are\n`calculator`.\n"}],"role":"user"},"tools":[{"functionDeclarations":[{"description":"This tool can now optionally accept skill_id and rationale parameters to guide skill-based orchestration. Transfer the question to another agent.\nThis tool hands off control to another agent when it's more suitable to answer the user's question according to the agent's description.","name":"transfer_to_agent","parameters":{"properties":{"agent_name":{"description":"the agent name to transfer to","enum":["calculator"],"type":"STRING"},"rationale":{"description":"The reasoning behind selecting this agent and skill.","type":"STRING"},"skill_id":{"description":"The specific skill to be utilized by the agent.","type":"STRING"}},"required":["agent_name"],"type":"OBJECT"}}]}]}HTTP/2.0 200 OK Content-Type: application/json; charset=UTF-8 Date: Tue, 10 Mar 2026 16:41:17 GMT Server: scaffolding on HTTPServer2 @@ -57,14 +57,14 @@ X-Xss-Protection: 0 "modelVersion": "gemini-2.5-flash", "responseId": "rEmwaZamPOmakdUP45HlkQo" } -2247 1254 +2273 1256 POST https://generativelanguage.googleapis.com/v1beta/models/gemini-2.5-flash:generateContent HTTP/1.1 Host: generativelanguage.googleapis.com User-Agent: Go-http-client/1.1 -Content-Length: 2014 +Content-Length: 2040 Content-Type: application/json -{"contents":[{"parts":[{"text":"Can you add 2 and 2?"}],"role":"user"},{"parts":[{"text":"For context:"},{"text":"[transfer_agent] called tool `transfer_to_agent` with parameters: {\"agent_name\":\"calculator\"}"}],"role":"user"},{"parts":[{"text":"For context:"},{"text":"[transfer_agent] `transfer_to_agent` tool returned result: {}"}],"role":"user"}],"generationConfig":{},"systemInstruction":{"parts":[{"text":"You are a calculator agent. You can calculate numbers.\n\nYou are an agent. Your internal name is \"calculator\". The description about you is \"calculator agent\\n Skills: add, subtract, multiply, divide\".\n\n\nYou have a list of other agents to transfer to:\n\n\nAgent name: transfer_agent\nAgent description: transfer agent\n\n\nIf you are the best to answer the question according to your description,\nyou can answer it.\n\nIf another agent is better for answering the question according to its\ndescription, call `transfer_to_agent` function to transfer the question to that\nagent. When transferring, do not generate any text other than the function\ncall.\n\n**NOTE**: the only available agents for `transfer_to_agent` function are\n`transfer_agent`.\n\nIf neither you nor the other agents are best for the question, transfer to your parent agent transfer_agent.\n"}],"role":"user"},"tools":[{"functionDeclarations":[{"description":"This tool can now optionally accept skill_id and rationale parameters to guide skill-based orchestration. Transfer the question to another agent.\nThis tool hands off control to another agent when it's more suitable to answer the user's question according to the agent's description.","name":"transfer_to_agent","parameters":{"properties":{"agent_name":{"description":"the agent name to transfer to","type":"string"},"rationale":{"description":"The reasoning behind selecting this agent and skill.","type":"STRING"},"skill_id":{"description":"The specific skill to be utilized by the agent.","type":"STRING"}},"required":["agent_name"],"type":"object"}}]}]}HTTP/2.0 200 OK +{"contents":[{"parts":[{"text":"Can you add 2 and 2?"}],"role":"user"},{"parts":[{"text":"For context:"},{"text":"[transfer_agent] called tool `transfer_to_agent` with parameters: {\"agent_name\":\"calculator\"}"}],"role":"user"},{"parts":[{"text":"For context:"},{"text":"[transfer_agent] `transfer_to_agent` tool returned result: {}"}],"role":"user"}],"generationConfig":{},"systemInstruction":{"parts":[{"text":"You are a calculator agent. You can calculate numbers.\n\nYou are an agent. Your internal name is \"calculator\". The description about you is \"calculator agent\\n Skills: add, subtract, multiply, divide\".\n\n\nYou have a list of other agents to transfer to:\n\n\nAgent name: transfer_agent\nAgent description: transfer agent\n\n\nIf you are the best to answer the question according to your description,\nyou can answer it.\n\nIf another agent is better for answering the question according to its\ndescription, call `transfer_to_agent` function to transfer the question to that\nagent. When transferring, do not generate any text other than the function\ncall.\n\n**NOTE**: the only available agents for `transfer_to_agent` function are\n`transfer_agent`.\n\nIf neither you nor the other agents are best for the question, transfer to your parent agent transfer_agent.\n"}],"role":"user"},"tools":[{"functionDeclarations":[{"description":"This tool can now optionally accept skill_id and rationale parameters to guide skill-based orchestration. Transfer the question to another agent.\nThis tool hands off control to another agent when it's more suitable to answer the user's question according to the agent's description.","name":"transfer_to_agent","parameters":{"properties":{"agent_name":{"description":"the agent name to transfer to","enum":["transfer_agent"],"type":"STRING"},"rationale":{"description":"The reasoning behind selecting this agent and skill.","type":"STRING"},"skill_id":{"description":"The specific skill to be utilized by the agent.","type":"STRING"}},"required":["agent_name"],"type":"OBJECT"}}]}]}HTTP/2.0 200 OK Content-Type: application/json; charset=UTF-8 Date: Tue, 10 Mar 2026 16:41:18 GMT Server: scaffolding on HTTPServer2 @@ -106,4 +106,4 @@ X-Xss-Protection: 0 }, "modelVersion": "gemini-2.5-flash", "responseId": "rUmwaZLzNJaxkdUP58PwoAk" -} + } diff --git a/plugin/loggingplugin/logging_plugin.go b/plugin/loggingplugin/logging_plugin.go index 83613a747..bd9e52aad 100644 --- a/plugin/loggingplugin/logging_plugin.go +++ b/plugin/loggingplugin/logging_plugin.go @@ -279,7 +279,7 @@ func (p *loggingPlugin) onModelError(ctx agent.CallbackContext, req *model.LLMRe return nil, nil } -func (p *loggingPlugin) beforeTool(ctx tool.Context, t tool.Tool, args map[string]any) (map[string]any, error) { +func (p *loggingPlugin) beforeTool(ctx agent.ToolContext, t tool.Tool, args map[string]any) (map[string]any, error) { p.log("🔧 TOOL STARTING") p.log(fmt.Sprintf(" Tool Name: %s", t.Name())) p.log(fmt.Sprintf(" Agent: %s", ctx.AgentName())) @@ -288,7 +288,7 @@ func (p *loggingPlugin) beforeTool(ctx tool.Context, t tool.Tool, args map[strin return nil, nil } -func (p *loggingPlugin) afterTool(ctx tool.Context, t tool.Tool, args, result map[string]any, err error) (map[string]any, error) { +func (p *loggingPlugin) afterTool(ctx agent.ToolContext, t tool.Tool, args, result map[string]any, err error) (map[string]any, error) { p.log("🔧 TOOL COMPLETED") p.log(fmt.Sprintf(" Tool Name: %s", t.Name())) p.log(fmt.Sprintf(" Agent: %s", ctx.AgentName())) @@ -301,7 +301,7 @@ func (p *loggingPlugin) afterTool(ctx tool.Context, t tool.Tool, args, result ma return nil, nil } -func (p *loggingPlugin) onToolError(ctx tool.Context, t tool.Tool, args map[string]any, err error) (map[string]any, error) { +func (p *loggingPlugin) onToolError(ctx agent.ToolContext, t tool.Tool, args map[string]any, err error) (map[string]any, error) { p.log("🔧 TOOL ERROR") p.log(fmt.Sprintf(" Tool Name: %s", t.Name())) p.log(fmt.Sprintf(" Agent: %s", ctx.AgentName())) diff --git a/plugin/plugin_manager_test.go b/plugin/plugin_manager_test.go index c081bf6f3..81c256125 100644 --- a/plugin/plugin_manager_test.go +++ b/plugin/plugin_manager_test.go @@ -36,7 +36,7 @@ import ( type testCase struct { name string - tool func(tool.Context, map[string]any) (map[string]any, error) + tool func(agent.ToolContext, map[string]any) (map[string]any, error) args map[string]any beforeToolCallbacks []llmagent.BeforeToolCallback afterToolCallbacks []llmagent.AfterToolCallback @@ -51,7 +51,7 @@ func TestCallTool(t *testing.T) { testCases := []testCase{ { name: "tool runs successfully", - tool: func(ctx tool.Context, args map[string]any) (map[string]any, error) { + tool: func(ctx agent.ToolContext, args map[string]any) (map[string]any, error) { return map[string]any{"result": "success"}, nil }, args: map[string]any{"key": "value"}, @@ -60,7 +60,7 @@ func TestCallTool(t *testing.T) { }, { name: "tool error", - tool: func(ctx tool.Context, args map[string]any) (map[string]any, error) { + tool: func(ctx agent.ToolContext, args map[string]any) (map[string]any, error) { return nil, errors.New("tool error") }, args: map[string]any{"key": "value"}, @@ -68,15 +68,15 @@ func TestCallTool(t *testing.T) { }, { name: "before callback returns result", - tool: func(ctx tool.Context, args map[string]any) (map[string]any, error) { + tool: func(ctx agent.ToolContext, args map[string]any) (map[string]any, error) { t.Error("tool should not be called") return nil, nil }, beforeToolCallbacks: []llmagent.BeforeToolCallback{ - func(ctx tool.Context, tool tool.Tool, args map[string]any) (map[string]any, error) { + func(ctx agent.ToolContext, tool tool.Tool, args map[string]any) (map[string]any, error) { return map[string]any{"result": "intercepted"}, nil }, - func(ctx tool.Context, tool tool.Tool, args map[string]any) (map[string]any, error) { + func(ctx agent.ToolContext, tool tool.Tool, args map[string]any) (map[string]any, error) { return map[string]any{"result": "2nd callback should not be called"}, nil }, }, @@ -86,15 +86,15 @@ func TestCallTool(t *testing.T) { }, { name: "before callback returns error", - tool: func(ctx tool.Context, args map[string]any) (map[string]any, error) { + tool: func(ctx agent.ToolContext, args map[string]any) (map[string]any, error) { t.Error("tool should not be called") return nil, nil }, beforeToolCallbacks: []llmagent.BeforeToolCallback{ - func(ctx tool.Context, tool tool.Tool, args map[string]any) (map[string]any, error) { + func(ctx agent.ToolContext, tool tool.Tool, args map[string]any) (map[string]any, error) { return nil, errors.New("before callback error") }, - func(ctx tool.Context, tool tool.Tool, args map[string]any) (map[string]any, error) { + func(ctx agent.ToolContext, tool tool.Tool, args map[string]any) (map[string]any, error) { return nil, errors.New("unexpected error") }, }, @@ -103,11 +103,11 @@ func TestCallTool(t *testing.T) { }, { name: "after callback modifies result", - tool: func(ctx tool.Context, args map[string]any) (map[string]any, error) { + tool: func(ctx agent.ToolContext, args map[string]any) (map[string]any, error) { return map[string]any{"result": "original"}, nil }, afterToolCallbacks: []llmagent.AfterToolCallback{ - func(ctx tool.Context, tool tool.Tool, args, result map[string]any, err error) (map[string]any, error) { + func(ctx agent.ToolContext, tool tool.Tool, args, result map[string]any, err error) (map[string]any, error) { return map[string]any{"result": "modified"}, nil }, }, @@ -117,17 +117,17 @@ func TestCallTool(t *testing.T) { }, { name: "after callback handles error and are run in symmetrical order", - tool: func(ctx tool.Context, args map[string]any) (map[string]any, error) { + tool: func(ctx agent.ToolContext, args map[string]any) (map[string]any, error) { return nil, errors.New("tool error") }, afterToolCallbacks: []llmagent.AfterToolCallback{ - func(ctx tool.Context, tool tool.Tool, args, result map[string]any, err error) (map[string]any, error) { + func(ctx agent.ToolContext, tool tool.Tool, args, result map[string]any, err error) (map[string]any, error) { if err != nil { return map[string]any{"result": "error handled"}, nil } return nil, nil }, - func(ctx tool.Context, tool tool.Tool, args, result map[string]any, err error) (map[string]any, error) { + func(ctx agent.ToolContext, tool tool.Tool, args, result map[string]any, err error) (map[string]any, error) { t.Errorf("unexpected call to after tool callback") return map[string]any{"result": "unexpected output"}, nil }, @@ -137,14 +137,14 @@ func TestCallTool(t *testing.T) { }, { name: "after callback returns error and are run in symmetrical order", - tool: func(ctx tool.Context, args map[string]any) (map[string]any, error) { + tool: func(ctx agent.ToolContext, args map[string]any) (map[string]any, error) { return map[string]any{"result": "success"}, nil }, afterToolCallbacks: []llmagent.AfterToolCallback{ - func(ctx tool.Context, tool tool.Tool, args, result map[string]any, err error) (map[string]any, error) { + func(ctx agent.ToolContext, tool tool.Tool, args, result map[string]any, err error) (map[string]any, error) { return nil, errors.New("after callback error") }, - func(ctx tool.Context, tool tool.Tool, args, result map[string]any, err error) (map[string]any, error) { + func(ctx agent.ToolContext, tool tool.Tool, args, result map[string]any, err error) (map[string]any, error) { t.Errorf("unexpected call to after tool callback") return nil, errors.New("unexpected error") }, @@ -155,16 +155,16 @@ func TestCallTool(t *testing.T) { }, { name: "no-op callbacks return func results", - tool: func(ctx tool.Context, args map[string]any) (map[string]any, error) { + tool: func(ctx agent.ToolContext, args map[string]any) (map[string]any, error) { return map[string]any{"result": "success"}, nil }, beforeToolCallbacks: []llmagent.BeforeToolCallback{ - func(ctx tool.Context, tool tool.Tool, args map[string]any) (map[string]any, error) { + func(ctx agent.ToolContext, tool tool.Tool, args map[string]any) (map[string]any, error) { return nil, nil }, }, afterToolCallbacks: []llmagent.AfterToolCallback{ - func(ctx tool.Context, tool tool.Tool, args, result map[string]any, err error) (map[string]any, error) { + func(ctx agent.ToolContext, tool tool.Tool, args, result map[string]any, err error) (map[string]any, error) { return nil, nil }, }, @@ -173,17 +173,17 @@ func TestCallTool(t *testing.T) { }, { name: "before callback result passed to after callback", - tool: func(ctx tool.Context, args map[string]any) (map[string]any, error) { + tool: func(ctx agent.ToolContext, args map[string]any) (map[string]any, error) { t.Error("tool should not be called") return nil, nil }, beforeToolCallbacks: []llmagent.BeforeToolCallback{ - func(ctx tool.Context, tool tool.Tool, args map[string]any) (map[string]any, error) { + func(ctx agent.ToolContext, tool tool.Tool, args map[string]any) (map[string]any, error) { return map[string]any{"result": "from_before"}, nil }, }, afterToolCallbacks: []llmagent.AfterToolCallback{ - func(ctx tool.Context, tool tool.Tool, args, result map[string]any, err error) (map[string]any, error) { + func(ctx agent.ToolContext, tool tool.Tool, args, result map[string]any, err error) (map[string]any, error) { if val, ok := result["result"]; !ok || val != "from_before" { return nil, errors.New("unexpected result in after callback") } @@ -197,17 +197,17 @@ func TestCallTool(t *testing.T) { }, { name: "before callback error passed to after callback", - tool: func(ctx tool.Context, args map[string]any) (map[string]any, error) { + tool: func(ctx agent.ToolContext, args map[string]any) (map[string]any, error) { t.Error("tool should not be called") return nil, nil }, beforeToolCallbacks: []llmagent.BeforeToolCallback{ - func(ctx tool.Context, tool tool.Tool, args map[string]any) (map[string]any, error) { + func(ctx agent.ToolContext, tool tool.Tool, args map[string]any) (map[string]any, error) { return nil, errors.New("error_from_before") }, }, afterToolCallbacks: []llmagent.AfterToolCallback{ - func(ctx tool.Context, tool tool.Tool, args, result map[string]any, err error) (map[string]any, error) { + func(ctx agent.ToolContext, tool tool.Tool, args, result map[string]any, err error) (map[string]any, error) { if err == nil || err.Error() != "error_from_before" { return nil, errors.New("unexpected error in after callback") } @@ -220,17 +220,17 @@ func TestCallTool(t *testing.T) { }, { name: "before callback error passed to on tool error callback", - tool: func(ctx tool.Context, args map[string]any) (map[string]any, error) { + tool: func(ctx agent.ToolContext, args map[string]any) (map[string]any, error) { t.Error("tool should not be called") return nil, nil }, beforeToolCallbacks: []llmagent.BeforeToolCallback{ - func(ctx tool.Context, tool tool.Tool, args map[string]any) (map[string]any, error) { + func(ctx agent.ToolContext, tool tool.Tool, args map[string]any) (map[string]any, error) { return nil, errors.New("error_from_before") }, }, onToolErrorCallbacks: []llmagent.OnToolErrorCallback{ - func(ctx tool.Context, tool tool.Tool, args map[string]any, err error) (map[string]any, error) { + func(ctx agent.ToolContext, tool tool.Tool, args map[string]any, err error) (map[string]any, error) { if err == nil || err.Error() != "error_from_before" { return nil, errors.New("unexpected error in on tool error callback") } @@ -243,17 +243,17 @@ func TestCallTool(t *testing.T) { }, { name: "before callback error passed to on tool error callback and after tool called", - tool: func(ctx tool.Context, args map[string]any) (map[string]any, error) { + tool: func(ctx agent.ToolContext, args map[string]any) (map[string]any, error) { t.Error("tool should not be called") return nil, nil }, beforeToolCallbacks: []llmagent.BeforeToolCallback{ - func(ctx tool.Context, tool tool.Tool, args map[string]any) (map[string]any, error) { + func(ctx agent.ToolContext, tool tool.Tool, args map[string]any) (map[string]any, error) { return nil, errors.New("error_from_before") }, }, onToolErrorCallbacks: []llmagent.OnToolErrorCallback{ - func(ctx tool.Context, tool tool.Tool, args map[string]any, err error) (map[string]any, error) { + func(ctx agent.ToolContext, tool tool.Tool, args map[string]any, err error) (map[string]any, error) { if err == nil || err.Error() != "error_from_before" { return nil, errors.New("unexpected error in on tool error callback") } @@ -261,7 +261,7 @@ func TestCallTool(t *testing.T) { }, }, afterToolCallbacks: []llmagent.AfterToolCallback{ - func(ctx tool.Context, tool tool.Tool, args, result map[string]any, err error) (map[string]any, error) { + func(ctx agent.ToolContext, tool tool.Tool, args, result map[string]any, err error) (map[string]any, error) { if err != nil { return nil, errors.New("unexpected error in after callback") } @@ -275,17 +275,17 @@ func TestCallTool(t *testing.T) { }, { name: "before callback error passed to on tool error callback and passed to after tool called", - tool: func(ctx tool.Context, args map[string]any) (map[string]any, error) { + tool: func(ctx agent.ToolContext, args map[string]any) (map[string]any, error) { t.Error("tool should not be called") return nil, nil }, beforeToolCallbacks: []llmagent.BeforeToolCallback{ - func(ctx tool.Context, tool tool.Tool, args map[string]any) (map[string]any, error) { + func(ctx agent.ToolContext, tool tool.Tool, args map[string]any) (map[string]any, error) { return nil, errors.New("error_from_before") }, }, onToolErrorCallbacks: []llmagent.OnToolErrorCallback{ - func(ctx tool.Context, tool tool.Tool, args map[string]any, err error) (map[string]any, error) { + func(ctx agent.ToolContext, tool tool.Tool, args map[string]any, err error) (map[string]any, error) { if err == nil || err.Error() != "error_from_before" { return nil, errors.New("unexpected error in on tool error callback") } @@ -293,7 +293,7 @@ func TestCallTool(t *testing.T) { }, }, afterToolCallbacks: []llmagent.AfterToolCallback{ - func(ctx tool.Context, tool tool.Tool, args, result map[string]any, err error) (map[string]any, error) { + func(ctx agent.ToolContext, tool tool.Tool, args, result map[string]any, err error) (map[string]any, error) { if err == nil || err.Error() != "error_from_on_tool_error" { return nil, errors.New("unexpected error in after callback") } @@ -307,17 +307,17 @@ func TestCallTool(t *testing.T) { }, { name: "before callback error passed to on tool error callback and passed to after tool called and handled", - tool: func(ctx tool.Context, args map[string]any) (map[string]any, error) { + tool: func(ctx agent.ToolContext, args map[string]any) (map[string]any, error) { t.Error("tool should not be called") return nil, nil }, beforeToolCallbacks: []llmagent.BeforeToolCallback{ - func(ctx tool.Context, tool tool.Tool, args map[string]any) (map[string]any, error) { + func(ctx agent.ToolContext, tool tool.Tool, args map[string]any) (map[string]any, error) { return nil, errors.New("error_from_before") }, }, onToolErrorCallbacks: []llmagent.OnToolErrorCallback{ - func(ctx tool.Context, tool tool.Tool, args map[string]any, err error) (map[string]any, error) { + func(ctx agent.ToolContext, tool tool.Tool, args map[string]any, err error) (map[string]any, error) { if err == nil || err.Error() != "error_from_before" { return nil, errors.New("unexpected error in on tool error callback") } @@ -325,7 +325,7 @@ func TestCallTool(t *testing.T) { }, }, afterToolCallbacks: []llmagent.AfterToolCallback{ - func(ctx tool.Context, tool tool.Tool, args, result map[string]any, err error) (map[string]any, error) { + func(ctx agent.ToolContext, tool tool.Tool, args, result map[string]any, err error) (map[string]any, error) { if err == nil || err.Error() != "error_from_on_tool_error" { return nil, errors.New("unexpected error in after tool callback") } @@ -386,7 +386,7 @@ func TestCallTool(t *testing.T) { beforeToolCallbacksCalled := false afterToolCallbacksCalled := false onToolErrorCallbacks := []llmagent.OnToolErrorCallback{ - func(ctx tool.Context, tool tool.Tool, args map[string]any, err error) (map[string]any, error) { + func(ctx agent.ToolContext, tool tool.Tool, args map[string]any, err error) (map[string]any, error) { onToolErrorCallbacksCalled = true if tc.dontRunOnErrorCanonicalCallback { t.Error("on tool error should not be called") @@ -395,7 +395,7 @@ func TestCallTool(t *testing.T) { }, } beforeToolCallbacks := []llmagent.BeforeToolCallback{ - func(ctx tool.Context, tool tool.Tool, args map[string]any) (map[string]any, error) { + func(ctx agent.ToolContext, tool tool.Tool, args map[string]any) (map[string]any, error) { beforeToolCallbacksCalled = true if tc.dontRunBeforeCanonicalCallback { t.Error("before Tool Callback should not be called") @@ -404,7 +404,7 @@ func TestCallTool(t *testing.T) { }, } afterToolCallbacks := []llmagent.AfterToolCallback{ - func(ctx tool.Context, tool tool.Tool, args, result map[string]any, err error) (map[string]any, error) { + func(ctx agent.ToolContext, tool tool.Tool, args, result map[string]any, err error) (map[string]any, error) { afterToolCallbacksCalled = true if tc.dontRunAfterCanonicalCallback { t.Error("after Tool Callback should not be called") diff --git a/plugin/retryandreflect/plugin.go b/plugin/retryandreflect/plugin.go index 2fa5e4163..67dd1bdb6 100644 --- a/plugin/retryandreflect/plugin.go +++ b/plugin/retryandreflect/plugin.go @@ -28,6 +28,7 @@ import ( "sync" "text/template" + "google.golang.org/adk/agent" "google.golang.org/adk/plugin" "google.golang.org/adk/tool" @@ -124,7 +125,7 @@ func MustNew(opts ...PluginOption) *plugin.Plugin { return p } -func (r *retryAndReflect) afterTool(ctx tool.Context, tool tool.Tool, args, result map[string]any, err error) (map[string]any, error) { +func (r *retryAndReflect) afterTool(ctx agent.ToolContext, tool tool.Tool, args, result map[string]any, err error) (map[string]any, error) { if err == nil { isReflectResponse := false if rt, ok := result["response_type"].(string); ok && rt == reflectAndRetryResponseType { @@ -139,11 +140,11 @@ func (r *retryAndReflect) afterTool(ctx tool.Context, tool tool.Tool, args, resu return nil, nil } -func (r *retryAndReflect) onToolError(ctx tool.Context, tool tool.Tool, args map[string]any, err error) (map[string]any, error) { +func (r *retryAndReflect) onToolError(ctx agent.ToolContext, tool tool.Tool, args map[string]any, err error) (map[string]any, error) { return r.handleToolError(ctx, tool, args, err) } -func (r *retryAndReflect) handleToolError(ctx tool.Context, failedTool tool.Tool, args map[string]any, err error) (map[string]any, error) { +func (r *retryAndReflect) handleToolError(ctx agent.ToolContext, failedTool tool.Tool, args map[string]any, err error) (map[string]any, error) { // skip if the error is tool.ErrConfirmationRequired. if errors.Is(err, tool.ErrConfirmationRequired) || errors.Is(err, tool.ErrConfirmationRejected) { return nil, nil @@ -179,14 +180,14 @@ func (r *retryAndReflect) handleToolError(ctx tool.Context, failedTool tool.Tool return r.createToolRetryExceedMsg(failedTool, args, err), nil } -func (r *retryAndReflect) scopeKey(ctx tool.Context) string { +func (r *retryAndReflect) scopeKey(ctx agent.ToolContext) string { if r.scope == Global { return globalScopeKey } return ctx.InvocationID() } -func (r *retryAndReflect) resetFailuresForTool(ctx tool.Context, toolName string) { +func (r *retryAndReflect) resetFailuresForTool(ctx agent.ToolContext, toolName string) { scopeKey := r.scopeKey(ctx) r.mu.Lock() diff --git a/plugin/retryandreflect/plugin_test.go b/plugin/retryandreflect/plugin_test.go index 8fd4ec13d..e0605af6b 100644 --- a/plugin/retryandreflect/plugin_test.go +++ b/plugin/retryandreflect/plugin_test.go @@ -19,7 +19,7 @@ import ( "strings" "testing" - "google.golang.org/adk/tool" + "google.golang.org/adk/agent" ) type mockTool struct { @@ -31,7 +31,7 @@ func (m *mockTool) Description() string { return "" } func (m *mockTool) IsLongRunning() bool { return false } type mockContext struct { - tool.Context + agent.ToolContext invocationID string } diff --git a/runner/live_runner_test.go b/runner/live_runner_test.go new file mode 100644 index 000000000..e68fef161 --- /dev/null +++ b/runner/live_runner_test.go @@ -0,0 +1,301 @@ +// Copyright 2026 Google LLC +// +// Licensed under the Apache License, Version 2.0 (the "License"); +// you may not use this file except in compliance with the License. +// You may obtain a copy of the License at +// +// http://www.apache.org/licenses/LICENSE-2.0 +// +// Unless required by applicable law or agreed to in writing, software +// distributed under the License is distributed on an "AS IS" BASIS, +// WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. +// See the License for the specific language governing permissions and +// limitations under the License. + +package runner + +import ( + "context" + "iter" + "strings" + "testing" + + "google.golang.org/genai" + + "google.golang.org/adk/agent" + "google.golang.org/adk/plugin" + "google.golang.org/adk/session" +) + +type mockLiveAgent struct { + agent.Agent + runLiveFn func(ctx agent.InvocationContext) (agent.LiveSession, iter.Seq2[*session.Event, error], error) +} + +func (m *mockLiveAgent) RunLive(ctx agent.InvocationContext) (agent.LiveSession, iter.Seq2[*session.Event, error], error) { + return m.runLiveFn(ctx) +} + +type dummyLiveSession struct{} + +func (d *dummyLiveSession) Send(req agent.LiveRequest) error { return nil } +func (d *dummyLiveSession) Close() error { return nil } + +func TestRunner_RunLive_Callbacks(t *testing.T) { + ctx := context.Background() + appName, userID, sessionID := "testApp", "testUser", "testSession" + + sessionService := session.InMemoryService() + _, err := sessionService.Create(ctx, &session.CreateRequest{ + AppName: appName, + UserID: userID, + SessionID: sessionID, + }) + if err != nil { + t.Fatal(err) + } + + var beforeRunCalled, afterRunCalled bool + + p, err := plugin.New(plugin.Config{ + Name: "test_plugin", + BeforeRunCallback: func(ctx agent.InvocationContext) (*genai.Content, error) { + beforeRunCalled = true + return nil, nil + }, + AfterRunCallback: func(ctx agent.InvocationContext) { + afterRunCalled = true + }, + }) + if err != nil { + t.Fatal(err) + } + + testAgent := must(agent.New(agent.Config{Name: "test_agent"})) + mockLive := &mockLiveAgent{ + Agent: testAgent, + runLiveFn: func(ctx agent.InvocationContext) (agent.LiveSession, iter.Seq2[*session.Event, error], error) { + return &dummyLiveSession{}, func(yield func(*session.Event, error) bool) { + yield(session.NewEventWithContext(ctx, ctx.InvocationID()), nil) + }, nil + }, + } + + r, err := New(Config{ + AppName: appName, + Agent: mockLive, + SessionService: sessionService, + PluginConfig: PluginConfig{ + Plugins: []*plugin.Plugin{p}, + }, + }) + if err != nil { + t.Fatal(err) + } + + sess, iter, err := r.RunLive(ctx, userID, sessionID, agent.LiveRunConfig{}) + if err != nil { + t.Fatalf("RunLive failed: %v", err) + } + if sess == nil { + t.Fatal("expected LiveSession to be returned") + } + + if !beforeRunCalled { + t.Error("BeforeRunCallback was not called before starting RunLive") + } + + if afterRunCalled { + t.Error("AfterRunCallback should not be called until iterator is consumed") + } + + for _, err := range iter { + if err != nil { + t.Fatal(err) + } + } + + if !afterRunCalled { + t.Error("AfterRunCallback was not called after iterator was consumed") + } +} + +func TestRunner_RunLive_EarlyExit(t *testing.T) { + ctx := context.Background() + appName, userID, sessionID := "testApp", "testUser", "testSession2" + + sessionService := session.InMemoryService() + _, err := sessionService.Create(ctx, &session.CreateRequest{ + AppName: appName, + UserID: userID, + SessionID: sessionID, + }) + if err != nil { + t.Fatal(err) + } + + expectedContent := genai.NewContentFromText("early exit content", "") + + p, err := plugin.New(plugin.Config{ + Name: "test_plugin", + BeforeRunCallback: func(ctx agent.InvocationContext) (*genai.Content, error) { + return expectedContent, nil + }, + }) + if err != nil { + t.Fatal(err) + } + + testAgent := must(agent.New(agent.Config{Name: "test_agent"})) + var runLiveCalled bool + mockLive := &mockLiveAgent{ + Agent: testAgent, + runLiveFn: func(ctx agent.InvocationContext) (agent.LiveSession, iter.Seq2[*session.Event, error], error) { + runLiveCalled = true + return &dummyLiveSession{}, func(yield func(*session.Event, error) bool) {}, nil + }, + } + + r, err := New(Config{ + AppName: appName, + Agent: mockLive, + SessionService: sessionService, + PluginConfig: PluginConfig{ + Plugins: []*plugin.Plugin{p}, + }, + }) + if err != nil { + t.Fatal(err) + } + + sess, iter, err := r.RunLive(ctx, userID, sessionID, agent.LiveRunConfig{}) + if err != nil { + t.Fatalf("RunLive failed: %v", err) + } + if runLiveCalled { + t.Error("RunLive should not have been called on the agent due to early exit") + } + + var events []*session.Event + for ev, err := range iter { + if err != nil { + t.Fatal(err) + } + events = append(events, ev) + } + + if len(events) != 1 { + t.Fatalf("expected 1 event, got %d", len(events)) + } + + if events[0].LLMResponse.Content != expectedContent { + t.Errorf("expected content %v, got %v", expectedContent, events[0].LLMResponse.Content) + } + + err = sess.Send(agent.LiveRequest{}) + if err == nil || !strings.Contains(err.Error(), "session is closed") { + t.Errorf("expected error 'session is closed' when sending to early exited session, got %v", err) + } + if err := sess.Close(); err != nil { + t.Errorf("Close() failed: %v", err) + } +} + +func TestRunner_RunLive_ChronologicalBuffering(t *testing.T) { + ctx := context.Background() + appName, userID, sessionID := "testApp", "testUser", "testSession3" + + sessionService := session.InMemoryService() + _, err := sessionService.Create(ctx, &session.CreateRequest{ + AppName: appName, + UserID: userID, + SessionID: sessionID, + }) + if err != nil { + t.Fatal(err) + } + + testAgent := must(agent.New(agent.Config{Name: "test_agent"})) + mockLive := &mockLiveAgent{ + Agent: testAgent, + runLiveFn: func(ctx agent.InvocationContext) (agent.LiveSession, iter.Seq2[*session.Event, error], error) { + return &dummyLiveSession{}, func(yield func(*session.Event, error) bool) { + // 1. Partial Transcription + ev1 := session.NewEventWithContext(ctx, ctx.InvocationID()) + ev1.LLMResponse.Partial = true + ev1.LLMResponse.OutputTranscription = &genai.Transcription{Text: "Hello"} + if !yield(ev1, nil) { + return + } + + // 2. Function Call (happening during transcription) + ev2 := session.NewEventWithContext(ctx, ctx.InvocationID()) + ev2.LLMResponse.Content = &genai.Content{ + Parts: []*genai.Part{{FunctionCall: &genai.FunctionCall{Name: "test_func"}}}, + } + if !yield(ev2, nil) { + return + } + + // 3. Final Transcription + ev3 := session.NewEventWithContext(ctx, ctx.InvocationID()) + ev3.LLMResponse.OutputTranscription = &genai.Transcription{Text: "Hello there."} + if !yield(ev3, nil) { + return + } + }, nil + }, + } + + r, err := New(Config{ + AppName: appName, + Agent: mockLive, + SessionService: sessionService, + }) + if err != nil { + t.Fatal(err) + } + + _, iter, err := r.RunLive(ctx, userID, sessionID, agent.LiveRunConfig{}) + if err != nil { + t.Fatalf("RunLive failed: %v", err) + } + + // Consume iterator to execute everything + for _, err := range iter { + if err != nil { + t.Fatal(err) + } + } + + // Verify Session History + getResp, err := sessionService.Get(ctx, &session.GetRequest{ + AppName: appName, + UserID: userID, + SessionID: sessionID, + }) + if err != nil { + t.Fatal(err) + } + + events := getResp.Session.Events() + // We expect 2 saved events: Final Transcription first, Function Call second. + // (Partial Transcription is not saved). + if events.Len() != 2 { + t.Fatalf("expected 2 saved events in session, got %d", events.Len()) + } + + // First saved event should be the final transcription + if events.At(0).LLMResponse.OutputTranscription == nil { + t.Errorf("expected first saved event to be transcription, but got %v", events.At(0)) + } + + if events.At(0).LLMResponse.OutputTranscription.Text != "Hello there." { + t.Errorf("expected first saved event to be transcription with text: %q, got: %q", "Hello there.", events.At(0).LLMResponse.OutputTranscription.Text) + } + + // Second saved event should be the function call + if events.At(1).LLMResponse.Content == nil || events.At(1).LLMResponse.Content.Parts[0].FunctionCall == nil { + t.Errorf("expected second saved event to be function call, but got %v", events.At(1)) + } +} diff --git a/runner/runner.go b/runner/runner.go index d7afbe244..53b3e3cd0 100644 --- a/runner/runner.go +++ b/runner/runner.go @@ -227,7 +227,7 @@ func (r *Runner) Run(ctx context.Context, userID, sessionID string, msg *genai.C earlyExitResult, err := pluginManager.RunBeforeRunCallback(ctx) if earlyExitResult != nil || err != nil { - earlyExitEvent := session.NewEvent(ctx.InvocationID()) + earlyExitEvent := session.NewEventWithContext(ctx, ctx.InvocationID()) earlyExitEvent.Author = "user" earlyExitEvent.LLMResponse = model.LLMResponse{ Content: msg, @@ -277,6 +277,285 @@ func (r *Runner) Run(ctx context.Context, userID, sessionID string, msg *genai.C } } +type liveAgent interface { + RunLive(ctx agent.InvocationContext) (agent.LiveSession, iter.Seq2[*session.Event, error], error) +} + +// RunLive runs a live session for the agent, supporting bidirectional streaming. +type runnerLiveSession struct { + sess agent.LiveSession + r *Runner + iCtx agent.InvocationContext + storedSession session.Session +} + +func (s *runnerLiveSession) Send(req agent.LiveRequest) error { + err := s.sess.Send(req) + if err != nil { + return err + } + + // Save user text content to session history + if req.Content != nil && len(req.Content.Parts) > 0 { + // Skip function responses - they are handled separately + isFunctionResponse := false + for _, part := range req.Content.Parts { + if part.FunctionResponse != nil { + isFunctionResponse = true + break + } + } + + if !isFunctionResponse { + event := session.NewEventWithContext(s.iCtx, s.iCtx.InvocationID()) + event.Author = "user" + event.LLMResponse = model.LLMResponse{ + Content: req.Content, + } + if err := s.r.sessionService.AppendEvent(s.iCtx, s.storedSession, event); err != nil { + return fmt.Errorf("failed to add user event to session: %w", err) + } + } + } + + return nil +} + +func (s *runnerLiveSession) Close() error { + return s.sess.Close() +} + +type closedLiveSession struct{} + +func (s *closedLiveSession) Send(req agent.LiveRequest) error { + return fmt.Errorf("session is closed") +} + +func (s *closedLiveSession) Close() error { + return nil +} + +// RunLive runs the agent in live bidirectional streaming mode via the +// upstream (google/adk-go) live engine introduced in v1.5.0, returning an +// [agent.LiveSession] for sending client events alongside the event stream. +// +// Fork note: this is not the method named RunLive before the v1.5.0 merge — +// the fork's queue-based entry point is now [Runner.RunLiveQueue] and remains +// the hulilabs-supported live path. The upstream engine behind this method +// has known gaps (see [Runner.RunLiveQueue] and the +// internal/llminternal/liveflow package doc) and its events never carry +// [session.LiveDiagnostics]. Pre-v1.5.0 RunLive callers should migrate to +// RunLiveQueue, not to this method. +func (r *Runner) RunLive(ctx context.Context, userID, sessionID string, cfg agent.LiveRunConfig, opts ...RunOption) (agent.LiveSession, iter.Seq2[*session.Event, error], error) { + options := runOptions{} + for _, opt := range opts { + opt(&options) + } + + var storedSession session.Session + getResp, err := r.sessionService.Get(ctx, &session.GetRequest{ + AppName: r.appName, + UserID: userID, + SessionID: sessionID, + }) + if err != nil { + if !r.autoCreateSession { + return nil, nil, err + } + createResp, err := r.sessionService.Create(ctx, &session.CreateRequest{ + AppName: r.appName, + UserID: userID, + SessionID: sessionID, + }) + if err != nil { + return nil, nil, err + } + storedSession = createResp.Session + } else { + storedSession = getResp.Session + } + + // msg is nil for Live run as it's streaming + agentToRun, err := r.findAgentToRun(storedSession, nil) + if err != nil { + return nil, nil, err + } + + lAgent, ok := agentToRun.(liveAgent) + if !ok { + return nil, nil, fmt.Errorf("agent %s does not support Live Run", agentToRun.Name()) + } + + ctx = parentmap.ToContext(ctx, r.parents) + ctx = runconfig.ToContext(ctx, &runconfig.RunConfig{ + StreamingMode: runconfig.StreamingModeBidi, // Live is always bidirectional streaming + Live: &cfg, + }) + // Same guard as Run() and RunLiveQueue(): a plugin-less (sub-)runner must + // inherit the parent's plugin manager from context instead of clobbering + // it with its own empty one, so model/tool/agent callbacks still propagate. + if r.pluginManager != nil && r.pluginManager.HasPlugins() { + ctx = plugininternal.ToContext(ctx, r.pluginManager) + } + + var artifacts agent.Artifacts + if r.artifactService != nil { + artifacts = &artifactinternal.Artifacts{ + Service: r.artifactService, + SessionID: storedSession.ID(), + AppName: storedSession.AppName(), + UserID: storedSession.UserID(), + } + } + + var memoryImpl agent.Memory = nil + if r.memoryService != nil { + memoryImpl = &imemory.Memory{ + Service: r.memoryService, + SessionID: storedSession.ID(), + UserID: storedSession.UserID(), + AppName: storedSession.AppName(), + } + } + + iCtx := icontext.NewInvocationContext(ctx, icontext.InvocationContextParams{ + Artifacts: artifacts, + Memory: memoryImpl, + Session: storedSession, + Agent: agentToRun, + UserContent: nil, + }) + + if r.pluginManager != nil { + earlyExitResult, err := r.pluginManager.RunBeforeRunCallback(iCtx) + if err != nil { + return nil, nil, err + } + if earlyExitResult != nil { + earlyExitEvent := session.NewEventWithContext(iCtx, iCtx.InvocationID()) + earlyExitEvent.Author = agentToRun.Name() + earlyExitEvent.LLMResponse = model.LLMResponse{ + Content: earlyExitResult, + } + if err := r.sessionService.AppendEvent(iCtx, storedSession, earlyExitEvent); err != nil { + return nil, nil, fmt.Errorf("failed to add event to session: %w", err) + } + + earlyExitIter := func(yield func(*session.Event, error) bool) { + yield(earlyExitEvent, nil) + } + return &closedLiveSession{}, earlyExitIter, nil + } + } + + agentSess, innerIter, err := lAgent.RunLive(iCtx) + if err != nil { + return nil, nil, err + } + + wrappedIter := func(yield func(*session.Event, error) bool) { + if r.pluginManager != nil { + defer r.pluginManager.RunAfterRunCallback(iCtx) + } + + var bufferedEvents []*session.Event + isTranscribing := false + + for event, err := range innerIter { + if err != nil { + if !yield(nil, err) { + return + } + continue + } + + if r.pluginManager != nil { + modifiedEvent, pluginErr := r.pluginManager.RunOnEventCallback(iCtx, event) + if pluginErr != nil { + if !yield(nil, pluginErr) { + return + } + continue + } + if modifiedEvent != nil { + event = modifiedEvent + } + } + + // Chronological event buffering logic for Live streaming. + // Holds back tool calls/responses if they arrive before the transcription finishes. + if event.LLMResponse.Partial && (event.LLMResponse.InputTranscription != nil || event.LLMResponse.OutputTranscription != nil) { + isTranscribing = true + } + + isToolCallOrResp := false + if event.LLMResponse.Content != nil { + for _, part := range event.LLMResponse.Content.Parts { + if part.FunctionCall != nil || part.FunctionResponse != nil { + isToolCallOrResp = true + break + } + } + } + + if isTranscribing && isToolCallOrResp { + bufferedEvents = append(bufferedEvents, event) + continue + } + + if !event.LLMResponse.Partial { + if event.LLMResponse.InputTranscription != nil || event.LLMResponse.OutputTranscription != nil { + isTranscribing = false + + if err := r.sessionService.AppendEvent(iCtx, storedSession, event); err != nil { + if !yield(nil, fmt.Errorf("failed to add event to session: %w", err)) { + return + } + continue + } + if !yield(event, nil) { + return + } + + for _, bufferedEvent := range bufferedEvents { + if err := r.sessionService.AppendEvent(iCtx, storedSession, bufferedEvent); err != nil { + if !yield(nil, fmt.Errorf("failed to add event to session: %w", err)) { + return + } + continue + } + if !yield(bufferedEvent, nil) { + return + } + } + bufferedEvents = nil + continue + } + } + + if !event.LLMResponse.Partial && !hasInlineData(event) { + if err := r.sessionService.AppendEvent(iCtx, storedSession, event); err != nil { + if !yield(nil, fmt.Errorf("failed to add event to session: %w", err)) { + return + } + continue + } + } + + if !yield(event, nil) { + return + } + } + } + + return &runnerLiveSession{ + sess: agentSess, + r: r, + iCtx: iCtx, + storedSession: storedSession, + }, wrappedIter, nil +} + func (r *Runner) appendMessageToSession(ctx agent.InvocationContext, storedSession session.Session, msg *genai.Content, saveInputBlobsAsArtifacts bool, pluginManager *plugininternal.PluginManager, stateDelta map[string]any) (agent.InvocationContext, error) { if msg == nil { return ctx, nil @@ -318,7 +597,7 @@ func (r *Runner) appendMessageToSession(ctx agent.InvocationContext, storedSessi } } - event := session.NewEvent(ctx.InvocationID()) + event := session.NewEventWithContext(ctx, ctx.InvocationID()) event.Author = "user" event.LLMResponse = model.LLMResponse{ @@ -334,10 +613,20 @@ func (r *Runner) appendMessageToSession(ctx agent.InvocationContext, storedSessi return ctx, nil } -// RunLive runs the agent in live bidirectional streaming mode. +// RunLiveQueue runs the agent in live bidirectional streaming mode. // Unlike Run(), it always uses the root agent and does not append an initial user message. // Audio events (CustomMetadata["is_audio"]==true) are not persisted to the session. -func (r *Runner) RunLive( +// +// RunLiveQueue is the hulilabs-supported live entry point: it drives the +// internal/llminternal/liveflow engine via a *agent.LiveRequestQueue and +// returns the event stream as an iterator. It coexists with the upstream +// [Runner.RunLive] (agent.LiveSession API) that arrived in v1.5.0, which +// drives the upstream live engine — as of v1.5.0 that engine has known +// gaps (string-matched GoAway, no reconnect backoff, history re-sent on +// resume, tool cancellation ignored; see the liveflow package doc for the +// full comparison). The method was renamed from RunLive so upstream keeps +// that name and future syncs stay conflict-free. +func (r *Runner) RunLiveQueue( ctx context.Context, userID, sessionID string, queue *agent.LiveRequestQueue, @@ -734,3 +1023,15 @@ func (r *Runner) isTransferableAcrossAgentTree(agentToRun agent.Agent) bool { func isToolEvent(event *session.Event) bool { return len(utils.FunctionCalls(event.Content)) > 0 || len(utils.FunctionResponses(event.Content)) > 0 } + +func hasInlineData(event *session.Event) bool { + if event.LLMResponse.Content == nil { + return false + } + for _, part := range event.LLMResponse.Content.Parts { + if part.InlineData != nil { + return true + } + } + return false +} diff --git a/runner/runner_live_test.go b/runner/runner_live_test.go index 2f080eef3..57a7eb840 100644 --- a/runner/runner_live_test.go +++ b/runner/runner_live_test.go @@ -304,7 +304,7 @@ func collectEvents(t *testing.T, r *Runner, queue *agent.LiveRequestQueue) ([]*s t.Helper() var events []*session.Event var errs []error - for ev, err := range r.RunLive(context.Background(), "user1", "sess1", queue, agent.RunConfig{}) { + for ev, err := range r.RunLiveQueue(context.Background(), "user1", "sess1", queue, agent.RunConfig{}) { if err != nil { errs = append(errs, err) } @@ -598,7 +598,7 @@ func TestScenario3_InFlightToolCancellation(t *testing.T) { start := time.Now() cfg := agent.RunConfig{ToolCoalesceWindow: 10 * time.Millisecond} var events []*session.Event - for ev, err := range r.RunLive(context.Background(), "user1", "sess1", queue, cfg) { + for ev, err := range r.RunLiveQueue(context.Background(), "user1", "sess1", queue, cfg) { if err != nil { break } @@ -1100,7 +1100,7 @@ func TestScenario10_ToolCallCoalescing(t *testing.T) { cfg := agent.RunConfig{ToolCoalesceWindow: 100 * time.Millisecond} var events []*session.Event - for ev, err := range r.RunLive(context.Background(), "user1", "sess1", queue, cfg) { + for ev, err := range r.RunLiveQueue(context.Background(), "user1", "sess1", queue, cfg) { if err != nil { break } @@ -1186,7 +1186,7 @@ func TestScenario11_ModelSpeakingStateTransitions(t *testing.T) { queue.Close() }() - for ev, err := range r.RunLive(context.Background(), "user1", "sess1", queue, agent.RunConfig{}) { + for ev, err := range r.RunLiveQueue(context.Background(), "user1", "sess1", queue, agent.RunConfig{}) { _ = ev _ = err } @@ -1391,7 +1391,7 @@ func TestScenario15_DeferFlushOnConsumerBreak(t *testing.T) { // Consumer collects 2 events then breaks (simulates user pressing Stop). collected := 0 - for ev, err := range r.RunLive(context.Background(), "user1", "sess1", queue, agent.RunConfig{}) { + for ev, err := range r.RunLiveQueue(context.Background(), "user1", "sess1", queue, agent.RunConfig{}) { _ = ev if err != nil { t.Fatalf("unexpected error: %v", err) @@ -1431,7 +1431,7 @@ func TestRunLive_LiveDiagnostics_NotLeakedToSession(t *testing.T) { queue.Close() var yieldedEvents []*session.Event - for ev, err := range r.RunLive(context.Background(), "user1", "sess1", queue, agent.RunConfig{}) { + for ev, err := range r.RunLiveQueue(context.Background(), "user1", "sess1", queue, agent.RunConfig{}) { if err != nil { t.Fatal(err) } @@ -1729,7 +1729,7 @@ func TestScenario21_LateOrphanFlushAfterTurnComplete(t *testing.T) { cfg := agent.RunConfig{ToolCoalesceWindow: 10 * time.Millisecond} var events []*session.Event - for ev, err := range r.RunLive(context.Background(), "user1", "sess1", queue, cfg) { + for ev, err := range r.RunLiveQueue(context.Background(), "user1", "sess1", queue, cfg) { if err != nil { break } @@ -2420,7 +2420,7 @@ func TestScenario33_EarlyExitFlushesBuffer(t *testing.T) { // Consumer reads 2 events (transcript + tool response from coalesce) then breaks. collected := 0 - for ev, err := range r.RunLive(context.Background(), "user1", "sess1", queue, agent.RunConfig{ToolCoalesceWindow: 10 * time.Millisecond}) { + for ev, err := range r.RunLiveQueue(context.Background(), "user1", "sess1", queue, agent.RunConfig{ToolCoalesceWindow: 10 * time.Millisecond}) { _ = ev if err != nil { t.Fatalf("unexpected error: %v", err) @@ -2508,7 +2508,7 @@ func TestScenario34_ReorderFlushBreakNoDoubleAppend(t *testing.T) { // remaining tail is re-flushed. cfg := agent.RunConfig{ToolCoalesceWindow: 5 * time.Millisecond} sawTool := false - for ev, err := range r.RunLive(context.Background(), "user1", "sess1", queue, cfg) { + for ev, err := range r.RunLiveQueue(context.Background(), "user1", "sess1", queue, cfg) { if err != nil { t.Fatalf("unexpected error: %v", err) } @@ -2791,7 +2791,7 @@ func TestScenario_NonThoughtCancelledThoughtOnlyDoesNotReset(t *testing.T) { cfg := agent.RunConfig{ToolCoalesceWindow: 500 * time.Millisecond} var events []*session.Event - for ev, err := range r.RunLive(context.Background(), "user1", "sess1", queue, cfg) { + for ev, err := range r.RunLiveQueue(context.Background(), "user1", "sess1", queue, cfg) { if err != nil { break } @@ -2984,7 +2984,7 @@ func TestScenario16_GoAwayReconnectionWithHandle(t *testing.T) { }() var events []*session.Event - for ev, err := range r.RunLive(context.Background(), "user1", "sess1", queue, agent.RunConfig{ + for ev, err := range r.RunLiveQueue(context.Background(), "user1", "sess1", queue, agent.RunConfig{ SessionResumption: &genai.SessionResumptionConfig{}, }) { if err != nil { @@ -3110,7 +3110,7 @@ func TestScenario17_NonResumableClearsHandle(t *testing.T) { queue.Close() }() - for ev, err := range r.RunLive(context.Background(), "user1", "sess1", queue, agent.RunConfig{ + for ev, err := range r.RunLiveQueue(context.Background(), "user1", "sess1", queue, agent.RunConfig{ SessionResumption: &genai.SessionResumptionConfig{}, }) { _ = ev @@ -3286,7 +3286,7 @@ func TestScenarioResumeUsesTransparentTrue(t *testing.T) { queue.Close() }() - for ev, err := range r.RunLive(context.Background(), "user1", "sess1", queue, agent.RunConfig{ + for ev, err := range r.RunLiveQueue(context.Background(), "user1", "sess1", queue, agent.RunConfig{ // Caller opts into session resumption with Transparent=false; the // resume path overwrites this on the second connect. SessionResumption: &genai.SessionResumptionConfig{Transparent: false}, @@ -3374,7 +3374,7 @@ func TestScenarioConnectionEOFWithHandleReconnects(t *testing.T) { }() var yieldedErr error - for ev, err := range r.RunLive(context.Background(), "user1", "sess1", queue, agent.RunConfig{ + for ev, err := range r.RunLiveQueue(context.Background(), "user1", "sess1", queue, agent.RunConfig{ SessionResumption: &genai.SessionResumptionConfig{}, }) { _ = ev @@ -3453,7 +3453,7 @@ func TestScenarioConnectionEOFWithoutHandleYieldsError(t *testing.T) { }() var yieldedErrs []error - for ev, err := range r.RunLive(context.Background(), "user1", "sess1", queue, agent.RunConfig{ + for ev, err := range r.RunLiveQueue(context.Background(), "user1", "sess1", queue, agent.RunConfig{ SessionResumption: &genai.SessionResumptionConfig{}, }) { _ = ev diff --git a/runner/runner_subagent_plugin_test.go b/runner/runner_subagent_plugin_test.go index 2bd859be5..f36c47254 100644 --- a/runner/runner_subagent_plugin_test.go +++ b/runner/runner_subagent_plugin_test.go @@ -307,7 +307,7 @@ func TestSubRunnerInheritsParentPlugins(t *testing.T) { queue := agent.NewLiveRequestQueue(100) queue.Close() - for _, err := range r.RunLive( + for _, err := range r.RunLiveQueue( context.Background(), "user1", "sess1", queue, agent.RunConfig{}, ) { if err != nil && err != io.EOF { @@ -412,7 +412,7 @@ func TestRunLiveInheritsContextPluginManager(t *testing.T) { queue := agent.NewLiveRequestQueue(100) queue.Close() - for _, err := range r.RunLive( + for _, err := range r.RunLiveQueue( seededCtx, "user1", "sess1", queue, agent.RunConfig{}, ) { if err != nil && err != io.EOF { @@ -428,3 +428,172 @@ func TestRunLiveInheritsContextPluginManager(t *testing.T) { t.Errorf("seeded plugin afterAgent = %d, want exactly 1 (RunLive inherited context manager)", afterAgent) } } + +// --------------------------------------------------------------------------- +// Regression test: the upstream-API RunLive entry point (agent.LiveSession +// signature, added in v1.5.0) must apply the same plugin-manager inheritance +// guard as Run and RunLiveQueue. +// +// The upstream live engine cannot run offline (it requires a *genai.Client +// live connection), so instead of driving it end-to-end this test uses a fake +// live agent that captures the plugin manager present in the invocation +// context RunLive hands to the agent. A plugin-less runner must leave the +// context-seeded parent manager in place; without the guard, the runner's own +// empty manager replaces it and parent plugin callbacks stop propagating. +// --------------------------------------------------------------------------- + +// noopLiveSession is a no-op agent.LiveSession returned by captureLiveAgent. +type noopLiveSession struct{} + +func (s *noopLiveSession) Send(agent.LiveRequest) error { return nil } +func (s *noopLiveSession) Close() error { return nil } + +// captureLiveAgent embeds a plain agent.Agent (the interface has an unexported +// method, so it cannot be implemented directly by a test type) and adds the +// RunLive method the runner's upstream live entry point type-asserts for, +// recording the plugin manager found in the invocation context. +type captureLiveAgent struct { + agent.Agent + mu sync.Mutex + captured *plugininternal.PluginManager +} + +func (a *captureLiveAgent) RunLive(ctx agent.InvocationContext) (agent.LiveSession, iter.Seq2[*session.Event, error], error) { + a.mu.Lock() + a.captured = plugininternal.FromContext(ctx) + a.mu.Unlock() + return &noopLiveSession{}, func(func(*session.Event, error) bool) {}, nil +} + +func (a *captureLiveAgent) capturedManager() *plugininternal.PluginManager { + a.mu.Lock() + defer a.mu.Unlock() + return a.captured +} + +func TestUpstreamRunLiveInheritsContextPluginManager(t *testing.T) { + counter := newCallCounter() + parentMgr, err := plugininternal.NewPluginManager(plugininternal.PluginConfig{ + Plugins: []*plugin.Plugin{counter.plugin(t)}, + }) + if err != nil { + t.Fatalf("plugininternal.NewPluginManager: %v", err) + } + seededCtx := plugininternal.ToContext(context.Background(), parentMgr) + + inner, err := agent.New(agent.Config{ + Name: "capture_agent", + Description: "captures the context plugin manager from RunLive", + Run: func(agent.InvocationContext) iter.Seq2[*session.Event, error] { + return func(func(*session.Event, error) bool) {} + }, + }) + if err != nil { + t.Fatalf("agent.New: %v", err) + } + capture := &captureLiveAgent{Agent: inner} + + svc := session.InMemoryService() + if _, err := svc.Create(context.Background(), &session.CreateRequest{ + AppName: "test", + UserID: "user1", + SessionID: "sess1", + }); err != nil { + t.Fatalf("session create: %v", err) + } + + // Empty PluginConfig: this runner owns no plugins, so its guard must leave + // the context-seeded parent manager in place. + r, err := runner.New(runner.Config{ + AppName: "test", + Agent: capture, + SessionService: svc, + PluginConfig: runner.PluginConfig{}, + }) + if err != nil { + t.Fatalf("runner.New: %v", err) + } + + sess, events, err := r.RunLive(seededCtx, "user1", "sess1", agent.LiveRunConfig{}) + if err != nil { + t.Fatalf("RunLive: %v", err) + } + defer func() { _ = sess.Close() }() + for _, err := range events { + if err != nil { + t.Fatalf("RunLive yielded error: %v", err) + } + } + + if got := capture.capturedManager(); got != parentMgr { + t.Errorf("RunLive installed plugin manager %p in the invocation context, want the inherited parent manager %p", got, parentMgr) + } +} + +// TestUpstreamRunLiveInstallsOwnPluginManager is the positive counterpart of +// TestUpstreamRunLiveInheritsContextPluginManager: a runner that owns plugins +// must install its own manager over a context-seeded parent one. Without this +// case, a guard regressed to never installing would still pass the suite. +func TestUpstreamRunLiveInstallsOwnPluginManager(t *testing.T) { + parentCounter := newCallCounter() + parentMgr, err := plugininternal.NewPluginManager(plugininternal.PluginConfig{ + Plugins: []*plugin.Plugin{parentCounter.plugin(t)}, + }) + if err != nil { + t.Fatalf("plugininternal.NewPluginManager: %v", err) + } + seededCtx := plugininternal.ToContext(context.Background(), parentMgr) + + inner, err := agent.New(agent.Config{ + Name: "capture_agent", + Description: "captures the context plugin manager from RunLive", + Run: func(agent.InvocationContext) iter.Seq2[*session.Event, error] { + return func(func(*session.Event, error) bool) {} + }, + }) + if err != nil { + t.Fatalf("agent.New: %v", err) + } + capture := &captureLiveAgent{Agent: inner} + + svc := session.InMemoryService() + if _, err := svc.Create(context.Background(), &session.CreateRequest{ + AppName: "test", + UserID: "user1", + SessionID: "sess1", + }); err != nil { + t.Fatalf("session create: %v", err) + } + + // This runner owns a plugin, so the HasPlugins guard must fire and replace + // the seeded parent manager with the runner's own. + ownCounter := newCallCounter() + r, err := runner.New(runner.Config{ + AppName: "test", + Agent: capture, + SessionService: svc, + PluginConfig: runner.PluginConfig{Plugins: []*plugin.Plugin{ownCounter.plugin(t)}}, + }) + if err != nil { + t.Fatalf("runner.New: %v", err) + } + + sess, events, err := r.RunLive(seededCtx, "user1", "sess1", agent.LiveRunConfig{}) + if err != nil { + t.Fatalf("RunLive: %v", err) + } + defer func() { _ = sess.Close() }() + for _, err := range events { + if err != nil { + t.Fatalf("RunLive yielded error: %v", err) + } + } + + got := capture.capturedManager() + if got == parentMgr { + t.Errorf("RunLive left the seeded parent manager %p in the invocation context; want the runner's own manager (guard must install when the runner has plugins)", parentMgr) + } + if !got.HasPlugins() { + t.Errorf("RunLive installed a plugin manager with no plugins in the invocation context; want the runner's own populated manager") + } +} diff --git a/server/adka2a/conversions.go b/server/adka2a/conversions.go new file mode 100644 index 000000000..b35e5ed97 --- /dev/null +++ b/server/adka2a/conversions.go @@ -0,0 +1,155 @@ +// Copyright 2025 Google LLC +// +// Licensed under the Apache License, Version 2.0 (the "License"); +// you may not use this file except in compliance with the License. +// You may obtain a copy of the License at +// +// http://www.apache.org/licenses/LICENSE-2.0 +// +// Unless required by applicable law or agreed to in writing, software +// distributed under the License is distributed on an "AS IS" BASIS, +// WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. +// See the License for the specific language governing permissions and +// limitations under the License. + +package adka2a + +import ( + "context" + "fmt" + + "github.com/a2aproject/a2a-go/a2a" + "github.com/a2aproject/a2a-go/log" + v2a2a "github.com/a2aproject/a2a-go/v2/a2a" + "github.com/a2aproject/a2a-go/v2/a2acompat/a2av0" + + "google.golang.org/genai" + + "google.golang.org/adk/agent" + v2 "google.golang.org/adk/server/adka2a/v2" + "google.golang.org/adk/session" +) + +// BuildAgentSkills attempts to create a list of [a2a.AgentSkill]s based on agent descriptions and types. +// This information can be used in [a2a.AgentCard] to help clients understand agent capabilities. +func BuildAgentSkills(agent agent.Agent) []a2a.AgentSkill { + v1Skills := v2.BuildAgentSkills(agent) + card := a2av0.FromV1AgentCard(&v2a2a.AgentCard{Skills: v1Skills}) + return card.Skills +} + +// ToA2AMetaKey adds a prefix used to differentiate ADK-related values stored in Metadata an A2A event. +func ToA2AMetaKey(key string) string { + return v2.ToA2AMetaKey(key) +} + +// ToADKMetaKey adds a prefix used to differentiate A2A-related values stored in custom metadata of an ADK session event. +func ToADKMetaKey(key string) string { + return v2.ToADKMetaKey(key) +} + +// ToCustomMetadata creates ADK custom metadata from A2A task and context IDs. +func ToCustomMetadata(taskID a2a.TaskID, contextID string) map[string]any { + return v2.ToCustomMetadata(v2a2a.TaskID(taskID), contextID) +} + +// GetA2ATaskInfo returns A2A task and context IDs if they are present in session event custom metadata. +func GetA2ATaskInfo(event *session.Event) (a2a.TaskID, string) { + taskID, contextID := v2.GetA2ATaskInfo(event) + return a2a.TaskID(taskID), contextID +} + +// NewRemoteAgentEvent create a new Event authored by the agent running in the provided invocation context. +func NewRemoteAgentEvent(ctx agent.InvocationContext) *session.Event { + return v2.NewRemoteAgentEvent(ctx) +} + +// EventToMessage converts the provided session event to A2A message. +func EventToMessage(event *session.Event) (*a2a.Message, error) { + v1Msg, err := v2.EventToMessage(event) + if err != nil { + return nil, err + } + return a2av0.FromV1Message(v1Msg), nil +} + +// ToSessionEvent converts the provided a2a event to session event authored by the agent running in the provided invocation context. +func ToSessionEvent(ctx agent.InvocationContext, event a2a.Event) (*session.Event, error) { + v1Event, err := a2av0.ToV1Event(event) + if err != nil { + return nil, fmt.Errorf("a2a event conversion failed: %w", err) + } + return v2.ToSessionEvent(ctx, v1Event) +} + +// IsPartial takes metadata of an A2A object (eg. a2a.Part, a2a.Artifact) and returns true if +// it was marked as partial based on the ADK partial flag set on the original ADK object. +func IsPartial(meta map[string]any) bool { + return v2.IsPartial(meta) +} + +// IsPartialFlagSet takes metadata of an A2A object (eg. a2a.Part, a2a.Artifact) and returns true if +// the ADK partial flag was set on it. +func IsPartialFlagSet(meta map[string]any) bool { + return v2.IsPartialFlagSet(meta) +} + +// ToA2APart converts the provided genai part to A2A equivalent. Long running tool IDs are used for attaching metadata to +// the relevant data parts. +func ToA2APart(part *genai.Part, longRunningToolIDs []string) (a2a.Part, error) { + v1p, err := v2.ToA2APart(part, longRunningToolIDs) + if err != nil { + return nil, err + } + return a2av0.FromV1Part(v1p), nil +} + +// ToA2AParts converts the provided genai parts to A2A equivalents. Long running tool IDs are used for attaching metadata to +// the relevant data parts. +func ToA2AParts(parts []*genai.Part, longRunningToolIDs []string) ([]a2a.Part, error) { + v1ps, err := v2.ToA2AParts(parts, longRunningToolIDs) + if err != nil { + return nil, err + } + result := make([]a2a.Part, len(v1ps)) + for i, v1p := range v1ps { + result[i] = a2av0.FromV1Part(v1p) + } + return result, nil +} + +// ToGenAIPart converts the provided A2A part to a genai equivalent. +func ToGenAIPart(part a2a.Part) (*genai.Part, error) { + return v2.ToGenAIPart(a2av0.ToV1Part(part)) +} + +// ToGenAIParts converts the provided A2A parts to genai equivalents. +func ToGenAIParts(parts []a2a.Part) ([]*genai.Part, error) { + v1ps := make([]*v2a2a.Part, len(parts)) + for i, p := range parts { + v1ps[i] = a2av0.ToV1Part(p) + } + return v2.ToGenAIParts(v1ps) +} + +// WithoutPartialArtifacts returns a slice of artifacts without partial artifacts. +// Partial artifacts are usually discarded (contain no parts) after agent invocation is finished. +func WithoutPartialArtifacts(artifacts []*a2a.Artifact) []*a2a.Artifact { + v1Artifacts := make([]*v2a2a.Artifact, 0, len(artifacts)) + for _, a := range artifacts { + v1a, err := a2av0.ToV1Artifact(a) + if err != nil { + log.Warn(context.Background(), "artifact skipped, because conversion failed", "cause", err) + continue + } + v1Artifacts = append(v1Artifacts, v1a) + } + + filtered := v2.WithoutPartialArtifacts(v1Artifacts) + + result := make([]*a2a.Artifact, 0, len(filtered)) + for _, f := range filtered { + result = append(result, a2av0.FromV1Artifact(f)) + } + return result +} diff --git a/server/adka2a/executor.go b/server/adka2a/executor.go index f17e344f0..990b7a2fb 100644 --- a/server/adka2a/executor.go +++ b/server/adka2a/executor.go @@ -12,27 +12,30 @@ // See the License for the specific language governing permissions and // limitations under the License. +// Package adka2a allows exposing ADK agents via A2A. +// +// Deprecated: Use google.golang.org/adk/server/adka2a/v2 instead. package adka2a import ( "context" "errors" "fmt" - "iter" - "slices" + "maps" "github.com/a2aproject/a2a-go/a2a" - "github.com/a2aproject/a2a-go/a2aclient" "github.com/a2aproject/a2a-go/a2asrv" "github.com/a2aproject/a2a-go/a2asrv/eventqueue" "github.com/a2aproject/a2a-go/log" - + a2av2 "github.com/a2aproject/a2a-go/v2/a2a" + "github.com/a2aproject/a2a-go/v2/a2acompat/a2av0" + a2asrvv2 "github.com/a2aproject/a2a-go/v2/a2asrv" "google.golang.org/genai" "google.golang.org/adk/agent" - iremoteagent "google.golang.org/adk/internal/agent/remoteagent" "google.golang.org/adk/plugin" "google.golang.org/adk/runner" + v2 "google.golang.org/adk/server/adka2a/v2" "google.golang.org/adk/session" ) @@ -55,18 +58,15 @@ type A2APartConverter func(ctx context.Context, a2aEvent a2a.Event, part a2a.Par // nil returns are considered intentionally dropped parts. type GenAIPartConverter func(ctx context.Context, adkEvent *session.Event, part *genai.Part) (a2a.Part, error) -// A2AExecutionCleanupCallback is a callback which will be called after an execution or cancellatio has completed or failed. +// A2AExecutionCleanupCallback is a callback which will be called after an execution or cancellation has completed or failed. type A2AExecutionCleanupCallback func(ctx context.Context, reqCtx *a2asrv.RequestContext, subAgentCards []*a2a.AgentCard, result a2a.SendMessageResult, cause error) // OutputMode controls how artifacts are produced. -type OutputMode string +type OutputMode = v2.OutputMode // Runner is an interface matching [runner.Runner] API. // It exists to let users use custom runner implementations with A2A agent executor. -type Runner interface { - // Run runs the agent for the given user input, yielding events from agents. - Run(ctx context.Context, userID, sessionID string, msg *genai.Content, cfg agent.RunConfig) iter.Seq2[*session.Event, error] -} +type Runner = v2.Runner // RunnerProvider is a [Runner] factory function. The provided plugin must be installed in the returned [Runner] for // callbacks taking [ExecutorContext] to work correctly. @@ -74,24 +74,17 @@ type RunnerProvider func(ctx context.Context, reqCtx *a2asrv.RequestContext, plu const ( // OutputArtifactPerRun produces a single artifact per [runner.Runner.Run]. - OutputArtifactPerRun OutputMode = "artifact-per-run" + OutputArtifactPerRun OutputMode = v2.OutputArtifactPerRun // OutputArtifactPerEvent produces an artifact per non-partial [session.Event]. // While agent is emitting events an artifact is build incrementally (parts are append to it). // The next partial event replaces accumulated contents and seals the artifact, meaning // the next event from this agent will create a new artifact. - OutputArtifactPerEvent OutputMode = "artifact-per-event" + OutputArtifactPerEvent OutputMode = v2.OutputArtifactPerEvent ) // RunnerConfig is part of the runner configuration executor code depends on. // Custom [RunnerProvider] needs to return it back to callers. -type RunnerConfig struct { - // AppName is the name of the application used in [session.Service] keys and A2A event metadata. - AppName string - // Agent is the root agent. It isued - Agent agent.Agent - // SessionService is the session service to use. - SessionService session.Service -} +type RunnerConfig = v2.RunnerConfig // ExecutorConfig allows to configure Executor. type ExecutorConfig struct { @@ -137,313 +130,279 @@ type ExecutorConfig struct { A2AExecutionCleanupCallback A2AExecutionCleanupCallback } -var _ a2asrv.AgentExecutor = (*Executor)(nil) - -// Executor invokes an ADK agent and translates [session.Event]-s to [a2a.Event]-s according to the following rules: -// - If the input doesn't reference any a2a.Task, produce a Task with TaskStateSubmitted state. -// - Right before runner.Runner invocation, produce TaskStatusUpdateEvent with TaskStateWorking. -// - For every session.Event produce a TaskArtifactUpdateEvent{Append=true} with transformed parts. -// - After the last session.Event is processed produce an empty TaskArtifactUpdateEvent{Append=true} with LastChunk=true, -// if at least one artifact update was produced during the run. -// - If there was an LLMResponse with non-zero error code, produce a TaskStatusUpdateEvent with TaskStateFailed. -// Else if there was an LLMResponse with long-running tool invocation, produce a TaskStatusUpdateEvent with TaskStateInputRequired. -// Else produce a TaskStatusUpdateEvent with TaskStateCompleted. +var ( + _ a2asrv.AgentExecutor = (*Executor)(nil) + _ a2asrv.AgentExecutionCleaner = (*Executor)(nil) +) + +// Executor is the legacy AgentExecutor implementation which delegates to [v2.Executor]. type Executor struct { - config ExecutorConfig + impl *v2.Executor } // NewExecutor creates an initialized [Executor] instance. func NewExecutor(config ExecutorConfig) *Executor { - if config.RunnerProvider == nil { - config.RunnerProvider = newDefaultRunnerProvider(config.RunnerConfig) + v1Config := v2.ExecutorConfig{ + RunnerConfig: config.RunnerConfig, + RunConfig: config.RunConfig, + OutputMode: v2.OutputMode(config.OutputMode), } - return &Executor{config: config} -} -func (e *Executor) Execute(ctx context.Context, reqCtx *a2asrv.RequestContext, queue eventqueue.Queue) error { - msg := reqCtx.Message - if msg == nil { - return fmt.Errorf("message not provided") + if config.RunnerProvider != nil { + v1Config.RunnerProvider = func(ctx context.Context, execCtx *a2asrvv2.ExecutorContext, plugin *plugin.Plugin) (RunnerConfig, Runner, error) { + return config.RunnerProvider(ctx, toRequestContext(execCtx), plugin) + } } - content, err := toGenAIContent(ctx, msg, e.config.A2APartConverter) - if err != nil { - return fmt.Errorf("a2a message conversion failed: %w", err) + + if config.BeforeExecuteCallback != nil { + v1Config.BeforeExecuteCallback = func(ctx context.Context, execCtx *a2asrvv2.ExecutorContext) (context.Context, error) { + legacyReqCtx := toRequestContext(execCtx) + newCtx, err := config.BeforeExecuteCallback(ctx, legacyReqCtx) + if err != nil { + return nil, err + } + v1ExecCtx, err := toExecutorContext(newCtx, legacyReqCtx) + if err != nil { + return nil, err + } + *execCtx = *v1ExecCtx + return newCtx, nil + } } - executorPlugin, err := newExecutorPlugin() - if err != nil { - return fmt.Errorf("failed to create a2a-executor plugin: %w", err) + if config.AfterEventCallback != nil { + v1Config.AfterEventCallback = func(ctx v2.ExecutorContext, adkEvent *session.Event, a2aEvent *a2av2.TaskArtifactUpdateEvent) error { + legacyEvent := a2av0.FromV1TaskArtifactUpdateEvent(a2aEvent) + if err := config.AfterEventCallback(executorContextWrapper{ctx}, adkEvent, legacyEvent); err != nil { + return err + } + newV1Event, err := a2av0.ToV1Event(legacyEvent) + if err != nil { + return err + } + if converted, ok := newV1Event.(*a2av2.TaskArtifactUpdateEvent); ok { + *a2aEvent = *converted + } + return nil + } } - cfg, r, err := e.config.RunnerProvider(ctx, reqCtx, executorPlugin.plugin) - if err != nil { - return fmt.Errorf("failed to create a runner: %w", err) + if config.AfterExecuteCallback != nil { + v1Config.AfterExecuteCallback = func(ctx v2.ExecutorContext, finalEvent *a2av2.TaskStatusUpdateEvent, err error) error { + legacyEvent := a2av0.FromV1TaskStatusUpdateEvent(finalEvent) + if cbErr := config.AfterExecuteCallback(executorContextWrapper{ctx}, legacyEvent, err); cbErr != nil { + return cbErr + } + newV1Event, convErr := a2av0.ToV1Event(legacyEvent) + if convErr != nil { + return convErr + } + if converted, ok := newV1Event.(*a2av2.TaskStatusUpdateEvent); ok { + *finalEvent = *converted + } + return nil + } } - if e.config.BeforeExecuteCallback != nil { - ctx, err = e.config.BeforeExecuteCallback(ctx, reqCtx) - if err != nil { - return fmt.Errorf("before execute: %w", err) + if config.A2APartConverter != nil { + v1Config.A2APartConverter = func(ctx context.Context, a2aEvent a2av2.Event, part *a2av2.Part) (*genai.Part, error) { + legacyEvent, err := a2av0.FromV1Event(a2aEvent) + if err != nil { + return nil, fmt.Errorf("a2a event conversion failed: %w", err) + } + return config.A2APartConverter(ctx, legacyEvent, a2av0.FromV1Part(part)) } } - if event, err := handleInputRequired(reqCtx, content); event != nil || err != nil { - if err != nil { - return err + if config.GenAIPartConverter != nil { + v1Config.GenAIPartConverter = func(ctx context.Context, adkEvent *session.Event, part *genai.Part) (*a2av2.Part, error) { + legacyPart, err := config.GenAIPartConverter(ctx, adkEvent, part) + if err != nil { + return nil, err + } + return a2av0.ToV1Part(legacyPart), nil } - return queue.Write(ctx, event) } - if reqCtx.StoredTask == nil { - event := a2a.NewSubmittedTask(reqCtx, msg) - if err := queue.Write(ctx, event); err != nil { - return fmt.Errorf("failed to submit a task: %w", err) + if config.A2AExecutionCleanupCallback != nil { + v1Config.A2AExecutionCleanupCallback = func(ctx context.Context, execCtx *a2asrvv2.ExecutorContext, subAgentCards []*a2av2.AgentCard, result a2av2.SendMessageResult, cause error) { + legacyReqCtx := toRequestContext(execCtx) + legacySubAgentCards := make([]*a2a.AgentCard, len(subAgentCards)) + for i, card := range subAgentCards { + legacySubAgentCards[i] = a2av0.FromV1AgentCard(card) + } + legacyEvent, err := a2av0.FromV1Event(result) + if err != nil { + log.Warn(ctx, "failed to convert SendMessageResult to legacy format", "error", err) + config.A2AExecutionCleanupCallback(ctx, legacyReqCtx, legacySubAgentCards, nil, errors.Join(cause, fmt.Errorf("SendMessageResult conversion failed: %w", err))) + return + } + legacyResult, ok := legacyEvent.(a2a.SendMessageResult) + if !ok { + log.Warn(ctx, "conversion result is not a2a.SendMessageResult", "type", fmt.Sprintf("%T", legacyResult)) + config.A2AExecutionCleanupCallback(ctx, legacyReqCtx, legacySubAgentCards, nil, errors.Join(cause, fmt.Errorf("conversion result is not a2a.SendMessageResult: %T", legacyEvent))) + return + } + config.A2AExecutionCleanupCallback(ctx, legacyReqCtx, legacySubAgentCards, legacyResult, cause) } } - invocationMeta := toInvocationMeta(ctx, cfg, reqCtx) + return &Executor{impl: v2.NewExecutor(v1Config)} +} - err = e.prepareSession(ctx, cfg, invocationMeta) +func (e *Executor) Execute(ctx context.Context, reqCtx *a2asrv.RequestContext, queue eventqueue.Queue) error { + execCtx, err := toExecutorContext(ctx, reqCtx) if err != nil { - event := toTaskFailedUpdateEvent(reqCtx, err, invocationMeta.eventMeta) - execCtx := newExecutorContext(ctx, invocationMeta, executorPlugin, content) - return e.writeFinalTaskStatus(execCtx, queue, nil, event, err) + return err } - event := a2a.NewStatusUpdateEvent(reqCtx, a2a.TaskStateWorking, nil) - event.Metadata = invocationMeta.eventMeta - if err := queue.Write(ctx, event); err != nil { - return err + ctx, callCtx := a2asrvv2.NewCallContext(ctx, execCtx.ServiceParams) + if execCtx.User.Authenticated { + callCtx.User = a2asrvv2.NewAuthenticatedUser(execCtx.User.Name, execCtx.User.Attributes) } - var artifactTransform eventToArtifactTransform - if e.config.OutputMode == OutputArtifactPerEvent { - artifactTransform = newArtifactMaker(reqCtx) - } else { - artifactTransform = newLegacyArtifactMaker(reqCtx) + for event, err := range e.impl.Execute(ctx, execCtx) { + if err != nil { + return err + } + legacyEvent, lErr := a2av0.FromV1Event(event) + if lErr != nil { + return lErr + } + if err := queue.Write(ctx, legacyEvent); err != nil { + return err + } } - processor := newEventProcessor(reqCtx, invocationMeta, e.config.GenAIPartConverter, artifactTransform) - executorContext := newExecutorContext(ctx, invocationMeta, executorPlugin, content) - return e.process(executorContext, r, processor, queue) + return nil } func (e *Executor) Cancel(ctx context.Context, reqCtx *a2asrv.RequestContext, queue eventqueue.Queue) error { - event := a2a.NewStatusUpdateEvent(reqCtx, a2a.TaskStateCanceled, nil) - event.Final = true - return queue.Write(ctx, event) -} - -func (e *Executor) Cleanup(ctx context.Context, reqCtx *a2asrv.RequestContext, result a2a.SendMessageResult, cause error) { - cfg, err := e.createRunnerConfig(ctx, reqCtx) + v2ReqCtx, err := toExecutorContext(ctx, reqCtx) if err != nil { - log.Error(ctx, "failed to create runner config", err) - return + return err } - remoteSubagents := findRemoteSubagents(cfg.Agent) - - // If task was in input-required and got successfully cancelled - run the cleanup logic - if reqCtx.StoredTask != nil && reqCtx.StoredTask.Status.State == a2a.TaskStateInputRequired { - if task, ok := result.(*a2a.Task); ok && task.Status.State == a2a.TaskStateCanceled && reqCtx.Message == nil { - if err := e.cancelChildInputRequiredTasks(ctx, reqCtx, reqCtx.StoredTask.Status, remoteSubagents); err != nil { - log.Warn(ctx, "failed to cancel subagent tasks waiting for input", "cause", err) - } + for event, err := range e.impl.Cancel(ctx, v2ReqCtx) { + if err != nil { + return err } - } - - if e.config.A2AExecutionCleanupCallback != nil { - subAgentCards := make([]*a2a.AgentCard, len(remoteSubagents)) - for i, subagent := range remoteSubagents { - subAgentCards[i] = subagent.config.AgentCard + legacyEvent, lErr := a2av0.FromV1Event(event) + if lErr != nil { + return lErr } - e.config.A2AExecutionCleanupCallback(ctx, reqCtx, subAgentCards, result, cause) - } else if cause != nil { - if reqCtx.Message != nil { - log.Warn(ctx, "execution failed", "error", cause) - } else { - log.Warn(ctx, "cancellation failed", "error", cause) + if err := queue.Write(ctx, legacyEvent); err != nil { + return err } } -} -func (e *Executor) cancelChildInputRequiredTasks(ctx context.Context, reqCtx *a2asrv.RequestContext, status a2a.TaskStatus, subagents []remoteAgent) error { - if len(subagents) == 0 { - return nil - } + return nil +} - cfg, err := e.createRunnerConfig(ctx, reqCtx) +func (e *Executor) Cleanup(ctx context.Context, reqCtx *a2asrv.RequestContext, result a2a.SendMessageResult, cause error) { + execCtx, err := toExecutorContext(ctx, reqCtx) if err != nil { - return fmt.Errorf("failed to create runner config: %w", err) + log.Warn(ctx, "failed to convert request context to executor context", "error", err) + return } - meta := toInvocationMeta(ctx, cfg, reqCtx) - getSessionResponse, err := cfg.SessionService.Get(ctx, &session.GetRequest{ - AppName: cfg.AppName, - UserID: meta.userID, - SessionID: meta.sessionID, - }) + v2Event, err := a2av0.ToV1Event(result) if err != nil { - return fmt.Errorf("failed to get a session: %w", err) + log.Warn(ctx, "failed to convert result to v2 event", "error", err) + return } - tasksToCancel, err := getSubagentTasksToCancel(ctx, status, getSessionResponse.Session) - if err != nil { - return fmt.Errorf("subtask search failed: %w", err) - } - if len(tasksToCancel) == 0 { - return nil + v2Result, ok := v2Event.(a2av2.SendMessageResult) + if !ok { + log.Warn(ctx, "converted event is not SendMessageResult", "type", fmt.Sprintf("%T", v2Event)) + return } - var failures []error - clientCache := map[string]*a2aclient.Client{} - for _, task := range tasksToCancel { // TODO(yarolegovich): run in parallel (how to limit?) - remoteSubagentIdx := slices.IndexFunc(subagents, func(a remoteAgent) bool { return a.agent.Name() == task.agentName }) - if remoteSubagentIdx < 0 { - continue - } - remoteSubagent := subagents[remoteSubagentIdx] - client, ok := clientCache[task.agentName] - if !ok { - _, newClient, err := iremoteagent.CreateA2AClient(ctx, remoteSubagent.config) - if err != nil { - failures = append(failures, fmt.Errorf("failed to create A2A client: %w", err)) - continue - } - clientCache[task.agentName] = newClient - client = newClient - } - _, err = client.CancelTask(ctx, &a2a.TaskIDParams{ID: task.taskID}) - if err != nil { - failures = append(failures, fmt.Errorf("failed to cancel task: %w", err)) - continue - } - } - for _, client := range clientCache { - if err := client.Destroy(); err != nil { - failures = append(failures, fmt.Errorf("client destroy failed: %w", err)) - } - } - return errors.Join(failures...) + e.impl.Cleanup(ctx, execCtx, v2Result, cause) } -// Processing failures should be delivered as Task failed events. An error is returned from this method if an event write fails. -func (e *Executor) process(ctx ExecutorContext, r Runner, processor *eventProcessor, q eventqueue.Queue) error { - meta := processor.meta - for adkEvent, adkErr := range r.Run(ctx, meta.userID, meta.sessionID, ctx.UserContent(), e.config.RunConfig) { - if adkErr != nil { - event := processor.makeTaskFailedEvent(fmt.Errorf("agent run failed: %w", adkErr), nil) - return e.writeFinalTaskStatus(ctx, q, processor.makeFinalArtifactUpdate(), event, adkErr) - } - - a2aEvent, pErr := processor.process(ctx, adkEvent) - if pErr == nil && a2aEvent != nil && e.config.AfterEventCallback != nil { - pErr = e.config.AfterEventCallback(ctx, adkEvent, a2aEvent) - } - - if pErr != nil { - event := processor.makeTaskFailedEvent(fmt.Errorf("processor failed: %w", pErr), adkEvent) - return e.writeFinalTaskStatus(ctx, q, processor.makeFinalArtifactUpdate(), event, pErr) - } - - if a2aEvent != nil { - if err := q.Write(ctx, a2aEvent); err != nil { - return fmt.Errorf("event write failed: %w", err) - } - } - } - - finalStatus := processor.makeFinalStatusUpdate() - return e.writeFinalTaskStatus(ctx, q, processor.makeFinalArtifactUpdate(), finalStatus, nil) +// ExecutorContext provides read-only information about the context of an A2A agent execution. +// An execution starts with a user message and ends with a task in a terminal or input-required state. +type ExecutorContext interface { + context.Context + + // SessionID is ID of the session. It is passed as contextID in A2A request. + SessionID() string + // UserID is ID of the user who made the request. The information is either extracted from [a2asrv.CallContext] + // or derived from session ID for unauthenticated requests. + UserID() string + // AgentName is the name of the root agent. + AgentName() string + // ReadonlyState provides a view of the current session state. + ReadonlyState() session.ReadonlyState + // Events provides a readonly view of the current session events. + Events() session.Events + // UserContent is a converted A2A message which is passed to runner.Run. + UserContent() *genai.Content + // RequestContext contains information about the original A2A Request, the current task and related tasks. + RequestContext() *a2asrv.RequestContext } -func (e *Executor) writeFinalTaskStatus( - ctx ExecutorContext, - queue eventqueue.Queue, - partialReset *a2a.TaskArtifactUpdateEvent, - status *a2a.TaskStatusUpdateEvent, - err error, -) error { - if e.config.AfterExecuteCallback != nil { - if err = e.config.AfterExecuteCallback(ctx, status, err); err != nil { - return fmt.Errorf("after execute: %w", err) - } - } - if partialReset != nil { - if err := queue.Write(ctx, partialReset); err != nil { - return fmt.Errorf("partial artifact update write failed: %w", err) - } - } - if err := queue.Write(ctx, status); err != nil { - return fmt.Errorf("%q state update event write failed: %w", status.Status.State, err) - } - return nil +type executorContextWrapper struct { + v2.ExecutorContext } -func (e *Executor) prepareSession(ctx context.Context, cfg RunnerConfig, meta invocationMeta) error { - service := cfg.SessionService +func (w executorContextWrapper) RequestContext() *a2asrv.RequestContext { + v1Ctx := w.ExecutorContext.RequestContext() + return toRequestContext(v1Ctx) +} - _, err := service.Get(ctx, &session.GetRequest{ - AppName: cfg.AppName, - UserID: meta.userID, - SessionID: meta.sessionID, - }) - if err == nil { - return nil +func toRequestContext(ctx *a2asrvv2.ExecutorContext) *a2asrv.RequestContext { + var relatedTasks []*a2a.Task + for _, t := range ctx.RelatedTasks { + relatedTasks = append(relatedTasks, a2av0.FromV1Task(t)) } - _, err = service.Create(ctx, &session.CreateRequest{ - AppName: cfg.AppName, - UserID: meta.userID, - SessionID: meta.sessionID, - State: make(map[string]any), - }) - if err != nil { - return fmt.Errorf("failed to create a session: %w", err) + return &a2asrv.RequestContext{ + ContextID: ctx.ContextID, + Message: a2av0.FromV1Message(ctx.Message), + StoredTask: a2av0.FromV1Task(ctx.StoredTask), + TaskID: a2a.TaskID(ctx.TaskID), + Metadata: ctx.Metadata, + RelatedTasks: relatedTasks, } - - return nil } -func (e *Executor) createRunnerConfig(ctx context.Context, reqCtx *a2asrv.RequestContext) (RunnerConfig, error) { - executorPlugin, err := newExecutorPlugin() - if err != nil { - return RunnerConfig{}, fmt.Errorf("failed to create a2a-plugin: %w", err) +func toExecutorContext(ctx context.Context, reqCtx *a2asrv.RequestContext) (*a2asrvv2.ExecutorContext, error) { + user := &a2asrvv2.User{Authenticated: false} + reqMeta := make(map[string][]string) + if callCtx, ok := a2asrv.CallContextFrom(ctx); ok { + user = &a2asrvv2.User{Name: callCtx.User.Name(), Authenticated: callCtx.User.Authenticated()} + maps.Insert(reqMeta, callCtx.RequestMeta().List()) } - cfg, _, err := e.config.RunnerProvider(ctx, reqCtx, executorPlugin.plugin) + + storedTask, err := a2av0.ToV1Task(reqCtx.StoredTask) if err != nil { - return RunnerConfig{}, fmt.Errorf("runner provider failed: %w", err) + return nil, err } - return cfg, nil -} - -func newDefaultRunnerProvider(baseConfig runner.Config) RunnerProvider { - return func(ctx context.Context, reqCtx *a2asrv.RequestContext, plugin *plugin.Plugin) (RunnerConfig, Runner, error) { - if baseConfig.Agent == nil { - return RunnerConfig{}, nil, fmt.Errorf("runner.Config.Agent is not provided") - } - if baseConfig.Agent == nil { - return RunnerConfig{}, nil, fmt.Errorf("runner.Config.SessionService is not provided") - } - cfg := baseConfig - cfg.PluginConfig.Plugins = append(slices.Clone(cfg.PluginConfig.Plugins), plugin) - r, err := runner.New(cfg) + var relatedTasks []*a2av2.Task + for _, t := range reqCtx.RelatedTasks { + v1Task, err := a2av0.ToV1Task(t) if err != nil { - return RunnerConfig{}, nil, err + return nil, err } - return toInternalRunnerConfig(cfg), &defaultRunner{runner: r}, nil + relatedTasks = append(relatedTasks, v1Task) } -} - -type defaultRunner struct { - runner *runner.Runner -} -func (r *defaultRunner) Run(ctx context.Context, userID, sessionID string, msg *genai.Content, cfg agent.RunConfig) iter.Seq2[*session.Event, error] { - return r.runner.Run(ctx, userID, sessionID, msg, cfg) -} + v1Msg, err := a2av0.ToV1Message(reqCtx.Message) + if err != nil { + return nil, err + } -func toInternalRunnerConfig(cfg runner.Config) RunnerConfig { - return RunnerConfig{Agent: cfg.Agent, AppName: cfg.AppName, SessionService: cfg.SessionService} + return &a2asrvv2.ExecutorContext{ + ContextID: reqCtx.ContextID, + Message: v1Msg, + TaskID: a2av2.TaskID(reqCtx.TaskID), + StoredTask: storedTask, + RelatedTasks: relatedTasks, + Metadata: reqCtx.Metadata, + User: user, + ServiceParams: a2asrvv2.NewServiceParams(reqMeta), + }, nil } diff --git a/server/adka2a/agent_card.go b/server/adka2a/v2/agent_card.go similarity index 99% rename from server/adka2a/agent_card.go rename to server/adka2a/v2/agent_card.go index 27a13a6b9..f9d82133f 100644 --- a/server/adka2a/agent_card.go +++ b/server/adka2a/v2/agent_card.go @@ -20,7 +20,7 @@ import ( "slices" "strings" - "github.com/a2aproject/a2a-go/a2a" + "github.com/a2aproject/a2a-go/v2/a2a" "google.golang.org/adk/agent" "google.golang.org/adk/agent/workflowagents/loopagent" diff --git a/server/adka2a/agent_card_test.go b/server/adka2a/v2/agent_card_test.go similarity index 99% rename from server/adka2a/agent_card_test.go rename to server/adka2a/v2/agent_card_test.go index 628f1e927..d6990615b 100644 --- a/server/adka2a/agent_card_test.go +++ b/server/adka2a/v2/agent_card_test.go @@ -17,7 +17,7 @@ package adka2a import ( "testing" - "github.com/a2aproject/a2a-go/a2a" + "github.com/a2aproject/a2a-go/v2/a2a" "github.com/google/go-cmp/cmp" "google.golang.org/adk/agent" diff --git a/server/adka2a/doc.go b/server/adka2a/v2/doc.go similarity index 91% rename from server/adka2a/doc.go rename to server/adka2a/v2/doc.go index eb0ce7cad..07495145e 100644 --- a/server/adka2a/doc.go +++ b/server/adka2a/v2/doc.go @@ -12,5 +12,5 @@ // See the License for the specific language governing permissions and // limitations under the License. -// Package adka2a allows to expose ADK agents via A2A. +// Package adka2a allows exposing ADK agents via A2A. package adka2a diff --git a/server/adka2a/events.go b/server/adka2a/v2/events.go similarity index 68% rename from server/adka2a/events.go rename to server/adka2a/v2/events.go index 828e38404..714cedd97 100644 --- a/server/adka2a/events.go +++ b/server/adka2a/v2/events.go @@ -19,7 +19,7 @@ import ( "fmt" "maps" - "github.com/a2aproject/a2a-go/a2a" + "github.com/a2aproject/a2a-go/v2/a2a" "google.golang.org/genai" "google.golang.org/adk/agent" @@ -28,7 +28,7 @@ import ( // NewRemoteAgentEvent create a new Event authored by the agent running in the provided invocation context. func NewRemoteAgentEvent(ctx agent.InvocationContext) *session.Event { - event := session.NewEvent(ctx.InvocationID()) + event := session.NewEventWithContext(ctx, ctx.InvocationID()) event.Author = ctx.Agent().Name() event.Branch = ctx.Branch() return event @@ -39,8 +39,11 @@ func EventToMessage(event *session.Event) (*a2a.Message, error) { if event == nil { return nil, nil } - - parts, err := ToA2AParts(event.Content.Parts, event.LongRunningToolIDs) + var eventParts []*genai.Part + if event.Content != nil { + eventParts = event.Content.Parts + } + parts, err := ToA2AParts(eventParts, event.LongRunningToolIDs) if err != nil { return nil, fmt.Errorf("part conversion failed: %w", err) } @@ -58,7 +61,13 @@ func EventToMessage(event *session.Event) (*a2a.Message, error) { msg := a2a.NewMessage(role, parts...) msg.Metadata = setActionsMeta(msg.Metadata, event.Actions) - maps.Copy(msg.Metadata, eventMeta) + if len(eventMeta) > 0 { + if msg.Metadata == nil { + msg.Metadata = maps.Clone(eventMeta) + } else { + maps.Copy(msg.Metadata, eventMeta) + } + } return msg, nil } @@ -70,7 +79,7 @@ func ToSessionEvent(ctx agent.InvocationContext, event a2a.Event) (*session.Even // ToSessionEventWithParts converts the provided a2a event to session event with custom part converter. func ToSessionEventWithParts(ctx agent.InvocationContext, event a2a.Event, partConverter A2APartConverter) (*session.Event, error) { if partConverter == nil { - partConverter = func(ctx context.Context, a2aEvent a2a.Event, part a2a.Part) (*genai.Part, error) { + partConverter = func(ctx context.Context, a2aEvent a2a.Event, part *a2a.Part) (*genai.Part, error) { return ToGenAIPart(part) } } @@ -96,7 +105,6 @@ func ToSessionEventWithParts(ctx agent.InvocationContext, event a2a.Event, partC if len(event.Content.Parts) == 0 { return nil, nil } - event.LongRunningToolIDs = getLongRunningToolIDs(v.Artifact.Parts, event.Content.Parts) if err := processA2AMeta(v, event); err != nil { return nil, fmt.Errorf("metadata processing failed: %w", err) } @@ -112,25 +120,30 @@ func ToSessionEventWithParts(ctx agent.InvocationContext, event a2a.Event, partC return event, nil case *a2a.TaskStatusUpdateEvent: - if v.Final { + if v.Status.State.Terminal() || v.Status.State == a2a.TaskStateInputRequired { return finalTaskStatusUpdateToEvent(ctx, v, partConverter) } if v.Status.Message == nil { return nil, nil } event, err := messageToEvent(ctx, v.Status.Message, partConverter) - event.TurnComplete = false if err != nil { return nil, fmt.Errorf("custom metadata conversion failed: %w", err) } + if event == nil { + return nil, nil + } + event.TurnComplete = false if len(event.Content.Parts) == 0 { return nil, nil } if err := processA2AMeta(v, event); err != nil { return nil, fmt.Errorf("metadata processing failed: %w", err) } - for _, part := range event.Content.Parts { - part.Thought = true + for _, p := range event.Content.Parts { + if p.Text != "" { + p.Thought = true + } } event.Partial = true return event, nil @@ -209,13 +222,14 @@ func artifactUpdateEventToEvent(ctx agent.InvocationContext, update *a2a.TaskArt return nil, nil } - parts, err := convertParts(ctx, update, update.Artifact.Parts, partConverter) + allParts, err := convertParts(ctx, update, update.Artifact.Parts, partConverter) if err != nil { return nil, fmt.Errorf("failed to convert artifact parts: %w", err) } event := NewRemoteAgentEvent(ctx) - event.Content = genai.NewContentFromParts(parts, genai.RoleModel) + event.Content = genai.NewContentFromParts(filterNilParts(allParts), genai.RoleModel) + event.LongRunningToolIDs = getLongRunningToolIDs(update.Artifact.Parts, allParts) return event, nil } @@ -227,32 +241,32 @@ func taskToEvent(ctx agent.InvocationContext, task *a2a.Task, partConverter A2AP var parts []*genai.Part var longRunningToolIDs []string for _, artifact := range task.Artifacts { - artifactParts, err := convertParts(ctx, task, artifact.Parts, partConverter) + allParts, err := convertParts(ctx, task, artifact.Parts, partConverter) if err != nil { return nil, fmt.Errorf("failed to convert artifact parts: %w", err) } - lrtIDs := getLongRunningToolIDs(artifact.Parts, artifactParts) + lrtIDs := getLongRunningToolIDs(artifact.Parts, allParts) - parts = append(parts, artifactParts...) + parts = append(parts, filterNilParts(allParts)...) longRunningToolIDs = append(longRunningToolIDs, lrtIDs...) } event := NewRemoteAgentEvent(ctx) - if task.Status.Message != nil { - msgParts, err := convertParts(ctx, task, task.Status.Message.Parts, partConverter) - if err != nil { - return nil, fmt.Errorf("failed to convert status message parts: %w", err) - } - lrtIDs := getLongRunningToolIDs(task.Status.Message.Parts, msgParts) + convertedStatusMsg, err := convertStatusMessage(ctx, task, task.Status, partConverter) + if err != nil { + return nil, err + } + parts = append(parts, convertedStatusMsg.parts...) + longRunningToolIDs = append(longRunningToolIDs, convertedStatusMsg.longRunningToolIDs...) - if task.Status.State == a2a.TaskStateFailed && len(msgParts) == 1 && msgParts[0].Text != "" { - event.ErrorMessage = msgParts[0].Text + if task.Status.State == a2a.TaskStateFailed { + if convertedStatusMsg.errorMessage != "" { + event.ErrorMessage = convertedStatusMsg.errorMessage } else { - parts = append(parts, msgParts...) + event.ErrorMessage = "a2a task failed" } - longRunningToolIDs = append(longRunningToolIDs, lrtIDs...) } isTerminal := task.Status.State.Terminal() || task.Status.State == a2a.TaskStateInputRequired @@ -279,37 +293,40 @@ func finalTaskStatusUpdateToEvent(ctx agent.InvocationContext, update *a2a.TaskS event := NewRemoteAgentEvent(ctx) - var parts []*genai.Part - var err error - if update.Status.Message != nil { - parts, err = convertParts(ctx, update, update.Status.Message.Parts, partConverter) - if err != nil { - return nil, fmt.Errorf("failed to convert status message parts: %w", err) - } + convertedStatusMsg, err := convertStatusMessage(ctx, update, update.Status, partConverter) + if err != nil { + return nil, err } - if update.Status.State == a2a.TaskStateFailed && len(parts) == 1 && parts[0].Text != "" { - event.ErrorMessage = parts[0].Text - } else if len(parts) > 0 { - event.Content = genai.NewContentFromParts(parts, genai.RoleModel) + if len(convertedStatusMsg.parts) > 0 { + event.Content = genai.NewContentFromParts(convertedStatusMsg.parts, genai.RoleModel) + } + event.LongRunningToolIDs = convertedStatusMsg.longRunningToolIDs + + if update.Status.State == a2a.TaskStateFailed { + if convertedStatusMsg.errorMessage != "" { + event.ErrorMessage = convertedStatusMsg.errorMessage + } else { + event.ErrorMessage = "a2a task failed" + } } + if err := processA2AMeta(update, event); err != nil { return nil, fmt.Errorf("metadata processing failed: %w", err) } - if update.Status.Message != nil { - event.LongRunningToolIDs = getLongRunningToolIDs(update.Status.Message.Parts, parts) - } event.TurnComplete = true return event, nil } -func getLongRunningToolIDs(parts []a2a.Part, converted []*genai.Part) []string { +func getLongRunningToolIDs(parts []*a2a.Part, converted []*genai.Part) []string { var ids []string for i, part := range parts { - dp, ok := part.(a2a.DataPart) - if !ok { + if part.Data() == nil { continue } - if longRunning, ok := dp.Metadata[a2aDataPartMetaLongRunningKey].(bool); ok && longRunning { + if longRunning, ok := part.Meta()[a2aDataPartMetaLongRunningKey].(bool); ok && longRunning { + if i >= len(converted) || converted[i] == nil { + continue + } fnCall := converted[i] if fnCall.FunctionCall == nil { // TODO(yarolegovich): log a warning @@ -339,16 +356,60 @@ func toEventActions(meta map[string]any) session.EventActions { return result } -func convertParts(ctx agent.InvocationContext, event a2a.Event, parts []a2a.Part, partConverter A2APartConverter) ([]*genai.Part, error) { - var genaiParts []*genai.Part - for _, part := range parts { +// convertParts converts A2A parts to GenAI parts using the provided converter. +// The returned slice preserves index alignment with the input: parts for which +// the converter returns nil are kept as nil entries. Use filterNilParts to get a +// dense slice suitable for content creation. +func convertParts(ctx agent.InvocationContext, event a2a.Event, parts []*a2a.Part, partConverter A2APartConverter) ([]*genai.Part, error) { + genaiParts := make([]*genai.Part, len(parts)) + for i, part := range parts { genaiPart, err := partConverter(ctx, event, part) if err != nil { return nil, fmt.Errorf("failed to convert part: %w", err) } - if genaiPart != nil { - genaiParts = append(genaiParts, genaiPart) - } + genaiParts[i] = genaiPart } return genaiParts, nil } + +func filterNilParts(parts []*genai.Part) []*genai.Part { + var result []*genai.Part + for _, p := range parts { + if p != nil { + result = append(result, p) + } + } + return result +} + +type convertedStatusMessage struct { + errorMessage string + parts []*genai.Part + longRunningToolIDs []string +} + +func convertStatusMessage(ctx agent.InvocationContext, event a2a.Event, status a2a.TaskStatus, partConverter A2APartConverter) (*convertedStatusMessage, error) { + if status.Message == nil { + return &convertedStatusMessage{}, nil + } + errMessage := "" + parts := status.Message.Parts + if status.State == a2a.TaskStateFailed && len(parts) > 0 && parts[0].Text() != "" { + errMessage = parts[0].Text() + + isErrMessage, _ := parts[0].Metadata[metadataIsErrMessageKey].(bool) + if isErrMessage { + parts = parts[1:] + } + } + + convertedParts, err := convertParts(ctx, event, parts, partConverter) + if err != nil { + return nil, fmt.Errorf("failed to convert status message parts: %w", err) + } + return &convertedStatusMessage{ + errorMessage: errMessage, + parts: filterNilParts(convertedParts), + longRunningToolIDs: getLongRunningToolIDs(parts, convertedParts), + }, nil +} diff --git a/server/adka2a/events_test.go b/server/adka2a/v2/events_test.go similarity index 67% rename from server/adka2a/events_test.go rename to server/adka2a/v2/events_test.go index e0c6e563b..3d72d639f 100644 --- a/server/adka2a/events_test.go +++ b/server/adka2a/v2/events_test.go @@ -18,7 +18,7 @@ import ( "context" "testing" - "github.com/a2aproject/a2a-go/a2a" + "github.com/a2aproject/a2a-go/v2/a2a" "github.com/google/go-cmp/cmp" "github.com/google/go-cmp/cmp/cmpopts" "google.golang.org/genai" @@ -46,7 +46,7 @@ func TestToSessionEvent(t *testing.T) { { name: "message", input: &a2a.Message{ - Parts: []a2a.Part{a2a.TextPart{Text: "foo"}}, + Parts: a2a.ContentParts{a2a.NewTextPart("foo")}, TaskID: taskID, ContextID: contextID, Metadata: map[string]any{ @@ -77,7 +77,7 @@ func TestToSessionEvent(t *testing.T) { { name: "nil values", input: &a2a.Message{ - Parts: []a2a.Part{a2a.TextPart{Text: "foo"}}, + Parts: a2a.ContentParts{a2a.NewTextPart("foo")}, TaskID: taskID, ContextID: contextID, Metadata: map[string]any{ @@ -101,6 +101,7 @@ func TestToSessionEvent(t *testing.T) { input: &a2a.Message{ TaskID: taskID, ContextID: contextID, + Parts: a2a.ContentParts{}, }, want: &session.Event{ LLMResponse: model.LLMResponse{ @@ -122,19 +123,20 @@ func TestToSessionEvent(t *testing.T) { Artifacts: []*a2a.Artifact{ { // long running key is ignored for non-input-required states ID: a2a.NewArtifactID(), - Parts: []a2a.Part{ - a2a.DataPart{ - Data: map[string]any{"id": "get_weather", "args": map[string]any{"city": "Warsaw"}, "name": "GetWeather"}, - Metadata: map[string]any{a2aDataPartMetaTypeKey: a2aDataPartTypeFunctionCall, a2aDataPartMetaLongRunningKey: true}, - }, + Parts: a2a.ContentParts{ + func() *a2a.Part { + p := a2a.NewDataPart(map[string]any{"id": "get_weather", "args": map[string]any{"city": "Warsaw"}, "name": "GetWeather"}) + p.Metadata = map[string]any{a2aDataPartMetaTypeKey: a2aDataPartTypeFunctionCall, a2aDataPartMetaLongRunningKey: true} + return p + }(), }, }, - {ID: a2a.NewArtifactID(), Parts: a2a.ContentParts{a2a.TextPart{Text: "foo"}}}, - {ID: a2a.NewArtifactID(), Parts: a2a.ContentParts{a2a.TextPart{Text: "bar"}}}, + {ID: a2a.NewArtifactID(), Parts: a2a.ContentParts{a2a.NewTextPart("foo")}}, + {ID: a2a.NewArtifactID(), Parts: a2a.ContentParts{a2a.NewTextPart("bar")}}, }, Status: a2a.TaskStatus{ State: a2a.TaskStateCompleted, - Message: a2a.NewMessage(a2a.MessageRoleAgent, a2a.TextPart{Text: "done"}), + Message: a2a.NewMessage(a2a.MessageRoleAgent, a2a.NewTextPart("done")), }, Metadata: map[string]any{ metadataGroundingKey: map[string]any{"sourceFlaggingUris": []any{map[string]any{"sourceId": "id1"}}}, @@ -207,11 +209,12 @@ func TestToSessionEvent(t *testing.T) { Artifacts: []*a2a.Artifact{ { ID: a2a.NewArtifactID(), - Parts: []a2a.Part{ - a2a.DataPart{ - Data: map[string]any{"id": "get_weather", "args": map[string]any{"city": "Warsaw"}, "name": "GetWeather"}, - Metadata: map[string]any{a2aDataPartMetaTypeKey: a2aDataPartTypeFunctionCall, a2aDataPartMetaLongRunningKey: true}, - }, + Parts: a2a.ContentParts{ + func() *a2a.Part { + p := a2a.NewDataPart(map[string]any{"id": "get_weather", "args": map[string]any{"city": "Warsaw"}, "name": "GetWeather"}) + p.Metadata = map[string]any{a2aDataPartMetaTypeKey: a2aDataPartTypeFunctionCall, a2aDataPartMetaLongRunningKey: true} + return p + }(), }, }, }, @@ -247,7 +250,7 @@ func TestToSessionEvent(t *testing.T) { TaskID: taskID, ContextID: contextID, Artifact: &a2a.Artifact{ - ID: a2a.NewArtifactID(), Parts: a2a.ContentParts{a2a.TextPart{Text: "foo"}, a2a.TextPart{Text: "bar"}}, + ID: a2a.NewArtifactID(), Parts: a2a.ContentParts{a2a.NewTextPart("foo"), a2a.NewTextPart("bar")}, }, Metadata: map[string]any{ metadataGroundingKey: map[string]any{"sourceFlaggingUris": []any{map[string]any{"sourceId": "id1"}}}, @@ -281,7 +284,7 @@ func TestToSessionEvent(t *testing.T) { ContextID: contextID, Artifact: &a2a.Artifact{ ID: a2a.NewArtifactID(), - Parts: []a2a.Part{}, + Parts: a2a.ContentParts{}, }, }, want: nil, @@ -293,11 +296,12 @@ func TestToSessionEvent(t *testing.T) { ContextID: contextID, Artifact: &a2a.Artifact{ ID: a2a.NewArtifactID(), - Parts: []a2a.Part{ - a2a.DataPart{ - Data: map[string]any{"id": "get_weather", "args": map[string]any{"city": "Warsaw"}, "name": "GetWeather"}, - Metadata: map[string]any{a2aDataPartMetaTypeKey: a2aDataPartTypeFunctionCall, a2aDataPartMetaLongRunningKey: true}, - }, + Parts: a2a.ContentParts{ + func() *a2a.Part { + p := a2a.NewDataPart(map[string]any{"id": "get_weather", "args": map[string]any{"city": "Warsaw"}, "name": "GetWeather"}) + p.Metadata = map[string]any{a2aDataPartMetaTypeKey: a2aDataPartTypeFunctionCall, a2aDataPartMetaLongRunningKey: true} + return p + }(), }, }, }, @@ -328,10 +332,10 @@ func TestToSessionEvent(t *testing.T) { input: &a2a.TaskStatusUpdateEvent{ TaskID: taskID, ContextID: contextID, - Final: true, Status: a2a.TaskStatus{ + State: a2a.TaskStateCompleted, Message: &a2a.Message{ - Parts: []a2a.Part{a2a.TextPart{Text: "foo"}}, + Parts: a2a.ContentParts{a2a.NewTextPart("foo")}, }, }, Metadata: map[string]any{ @@ -360,7 +364,7 @@ func TestToSessionEvent(t *testing.T) { }, { name: "final task status update without message", - input: &a2a.TaskStatusUpdateEvent{TaskID: taskID, ContextID: contextID, Final: true}, + input: &a2a.TaskStatusUpdateEvent{TaskID: taskID, ContextID: contextID, Status: a2a.TaskStatus{State: a2a.TaskStateCompleted}}, want: &session.Event{ LLMResponse: model.LLMResponse{ CustomMetadata: map[string]any{ @@ -379,9 +383,9 @@ func TestToSessionEvent(t *testing.T) { TaskID: taskID, ContextID: contextID, Status: a2a.TaskStatus{ - State: a2a.TaskStateCompleted, + State: a2a.TaskStateWorking, Message: &a2a.Message{ - Parts: []a2a.Part{a2a.TextPart{Text: "foo"}}, + Parts: a2a.ContentParts{a2a.NewTextPart("foo")}, }, }, }, @@ -408,14 +412,14 @@ func TestToSessionEvent(t *testing.T) { input: &a2a.TaskStatusUpdateEvent{ TaskID: taskID, ContextID: contextID, - Final: true, Status: a2a.TaskStatus{ State: a2a.TaskStateFailed, - Message: &a2a.Message{Parts: []a2a.Part{a2a.TextPart{Text: "failed with an error"}}}, + Message: &a2a.Message{Parts: a2a.ContentParts{a2a.NewTextPart("failed with an error")}}, }, }, want: &session.Event{ LLMResponse: model.LLMResponse{ + Content: genai.NewContentFromParts([]*genai.Part{genai.NewPartFromText("failed with an error")}, genai.RoleModel), ErrorMessage: "failed with an error", CustomMetadata: map[string]any{ customMetaTaskIDKey: string(taskID), @@ -435,10 +439,10 @@ func TestToSessionEvent(t *testing.T) { Artifacts: []*a2a.Artifact{ { ID: "artifact-1", - Parts: []a2a.Part{ - a2a.TextPart{Text: "Checking weather..."}, - a2a.DataPart{ - Data: map[string]any{"id": "tool_1", "name": "GetWeather", "args": map[string]any{"city": "London"}}, + Parts: []*a2a.Part{ + a2a.NewTextPart("Checking weather..."), + { + Content: a2a.Data{Value: map[string]any{"id": "tool_1", "name": "GetWeather", "args": map[string]any{"city": "London"}}}, Metadata: map[string]any{ a2aDataPartMetaTypeKey: a2aDataPartTypeFunctionCall, a2aDataPartMetaLongRunningKey: true, @@ -448,9 +452,9 @@ func TestToSessionEvent(t *testing.T) { }, { ID: "artifact-2", - Parts: []a2a.Part{ - a2a.DataPart{ - Data: map[string]any{"id": "tool_2", "name": "GetNews", "args": map[string]any{"topic": "tech"}}, + Parts: []*a2a.Part{ + { + Content: a2a.Data{Value: map[string]any{"id": "tool_2", "name": "GetNews", "args": map[string]any{"topic": "tech"}}}, Metadata: map[string]any{ a2aDataPartMetaTypeKey: a2aDataPartTypeFunctionCall, a2aDataPartMetaLongRunningKey: true, @@ -480,18 +484,120 @@ func TestToSessionEvent(t *testing.T) { }, }, { - name: "task with single-part text status", + name: "failed task with single-part text status", input: &a2a.Task{ ID: taskID, ContextID: contextID, Status: a2a.TaskStatus{ State: a2a.TaskStateFailed, - Message: &a2a.Message{Parts: []a2a.Part{a2a.TextPart{Text: "failed with an error"}}}, + Message: &a2a.Message{Parts: a2a.ContentParts{a2a.NewTextPart("failed with an error")}}, + }, + }, + want: &session.Event{ + LLMResponse: model.LLMResponse{ + Content: genai.NewContentFromParts([]*genai.Part{genai.NewPartFromText("failed with an error")}, genai.RoleModel), + ErrorMessage: "failed with an error", + CustomMetadata: map[string]any{customMetaTaskIDKey: string(taskID), customMetaContextIDKey: contextID}, + TurnComplete: true, + }, + Author: agentName, + Branch: branch, + }, + }, + { + name: "completed task with single-part text status", + input: &a2a.Task{ + ID: taskID, + ContextID: contextID, + Status: a2a.TaskStatus{ + State: a2a.TaskStateCompleted, + Message: &a2a.Message{Parts: a2a.ContentParts{a2a.NewTextPart("failed with an error")}}, + }, + }, + want: &session.Event{ + LLMResponse: model.LLMResponse{ + Content: genai.NewContentFromParts([]*genai.Part{genai.NewPartFromText("failed with an error")}, genai.RoleModel), + CustomMetadata: map[string]any{customMetaTaskIDKey: string(taskID), customMetaContextIDKey: contextID}, + TurnComplete: true, + }, + Author: agentName, + Branch: branch, + }, + }, + { + name: "failed task with a multi-part status", + input: &a2a.Task{ + ID: taskID, + ContextID: contextID, + Status: a2a.TaskStatus{ + State: a2a.TaskStateFailed, + Message: a2a.NewMessage(a2a.MessageRoleAgent, + a2a.NewTextPart("failed with an error"), + a2a.NewTextPart("extra details"), + ), + }, + }, + want: &session.Event{ + LLMResponse: model.LLMResponse{ + ErrorMessage: "failed with an error", + Content: genai.NewContentFromParts([]*genai.Part{ + genai.NewPartFromText("failed with an error"), + genai.NewPartFromText("extra details"), + }, genai.RoleModel), + CustomMetadata: map[string]any{customMetaTaskIDKey: string(taskID), customMetaContextIDKey: contextID}, + TurnComplete: true, + }, + Author: agentName, + Branch: branch, + }, + }, + { + name: "failed task status with a multi-part status", + input: &a2a.TaskStatusUpdateEvent{ + TaskID: taskID, + ContextID: contextID, + Status: a2a.TaskStatus{ + State: a2a.TaskStateFailed, + Message: a2a.NewMessage(a2a.MessageRoleAgent, + a2a.NewTextPart("failed with an error"), + a2a.NewDataPart(map[string]any{"structured": true}), + ), + }, + }, + want: &session.Event{ + LLMResponse: model.LLMResponse{ + ErrorMessage: "failed with an error", + Content: genai.NewContentFromParts([]*genai.Part{ + genai.NewPartFromText("failed with an error"), + {InlineData: &genai.Blob{ + Data: []byte(`{"structured":true}`), + MIMEType: "text/plain", + }}, + }, genai.RoleModel), + CustomMetadata: map[string]any{customMetaTaskIDKey: string(taskID), customMetaContextIDKey: contextID}, + TurnComplete: true, + }, + Author: agentName, + Branch: branch, + }, + }, + { + name: "failed task status metadataIsErrMessageKey not in content", + input: &a2a.TaskStatusUpdateEvent{ + TaskID: taskID, + ContextID: contextID, + Status: a2a.TaskStatus{ + State: a2a.TaskStateFailed, + Message: a2a.NewMessage(a2a.MessageRoleAgent, + &a2a.Part{Content: a2a.Text("failed with an error"), Metadata: map[string]any{metadataIsErrMessageKey: true}}, + a2a.NewTextPart("extra details"), + ), }, }, want: &session.Event{ LLMResponse: model.LLMResponse{ ErrorMessage: "failed with an error", + Content: genai.NewContentFromParts([]*genai.Part{genai.NewPartFromText("extra details")}, genai.RoleModel), CustomMetadata: map[string]any{customMetaTaskIDKey: string(taskID), customMetaContextIDKey: contextID}, TurnComplete: true, }, @@ -499,6 +605,37 @@ func TestToSessionEvent(t *testing.T) { Branch: branch, }, }, + { + name: "input-required task status with err message and long-running tool", + input: &a2a.TaskStatusUpdateEvent{ + TaskID: taskID, + ContextID: contextID, + Status: a2a.TaskStatus{ + State: a2a.TaskStateFailed, + Message: a2a.NewMessage(a2a.MessageRoleAgent, + &a2a.Part{Content: a2a.Text("requesting input"), Metadata: map[string]any{metadataIsErrMessageKey: true}}, + func() *a2a.Part { + p := a2a.NewDataPart(map[string]any{"id": "tool_lr", "name": "LongRunning", "args": map[string]any{}}) + p.Metadata = map[string]any{a2aDataPartMetaTypeKey: a2aDataPartTypeFunctionCall, a2aDataPartMetaLongRunningKey: true} + return p + }(), + ), + }, + }, + want: &session.Event{ + LLMResponse: model.LLMResponse{ + Content: genai.NewContentFromParts([]*genai.Part{ + {FunctionCall: &genai.FunctionCall{ID: "tool_lr", Name: "LongRunning", Args: map[string]any{}}}, + }, genai.RoleModel), + ErrorMessage: "requesting input", + CustomMetadata: map[string]any{customMetaTaskIDKey: string(taskID), customMetaContextIDKey: contextID}, + TurnComplete: true, + }, + LongRunningToolIDs: []string{"tool_lr"}, + Author: agentName, + Branch: branch, + }, + }, } ignoreFields := []cmp.Option{ @@ -528,11 +665,11 @@ func TestToSessionEventWithParts_NilResultFiltered(t *testing.T) { t.Fatalf("failed to create an agent: %v", err) } - keepPart := a2a.TextPart{Text: "KEEP"} - dropPart := a2a.TextPart{Text: "DROP"} + keepPart := a2a.NewTextPart("KEEP") + dropPart := a2a.NewTextPart("DROP") - filterConverter := func(ctx context.Context, ev a2a.Event, p a2a.Part) (*genai.Part, error) { - if tp, ok := p.(a2a.TextPart); ok && tp.Text == "DROP" { + filterConverter := func(ctx context.Context, ev a2a.Event, p *a2a.Part) (*genai.Part, error) { + if p.Text() == "DROP" { return nil, nil } return ToGenAIPart(p) @@ -547,23 +684,23 @@ func TestToSessionEventWithParts_NilResultFiltered(t *testing.T) { input: &a2a.Task{ ID: taskID, ContextID: contextID, - Artifacts: []*a2a.Artifact{{Parts: []a2a.Part{keepPart, dropPart}}}, + Artifacts: []*a2a.Artifact{{Parts: []*a2a.Part{keepPart, dropPart}}}, Status: a2a.TaskStatus{ State: a2a.TaskStateCompleted, - Message: &a2a.Message{Parts: []a2a.Part{keepPart, dropPart}}, + Message: &a2a.Message{Parts: []*a2a.Part{keepPart, dropPart}}, }, }, }, { name: "message event", input: &a2a.Message{ - Parts: []a2a.Part{keepPart, dropPart}, + Parts: []*a2a.Part{keepPart, dropPart}, }, }, { name: "artifact update event", input: &a2a.TaskArtifactUpdateEvent{ - Artifact: &a2a.Artifact{Parts: []a2a.Part{keepPart, dropPart}}, + Artifact: &a2a.Artifact{Parts: []*a2a.Part{keepPart, dropPart}}, }, }, { @@ -571,9 +708,9 @@ func TestToSessionEventWithParts_NilResultFiltered(t *testing.T) { input: &a2a.TaskStatusUpdateEvent{ TaskID: taskID, ContextID: contextID, - Final: true, Status: a2a.TaskStatus{ - Message: &a2a.Message{Parts: []a2a.Part{keepPart, dropPart}}, + State: a2a.TaskStateCompleted, + Message: &a2a.Message{Parts: []*a2a.Part{keepPart, dropPart}}, }, }, }, diff --git a/server/adka2a/v2/executor.go b/server/adka2a/v2/executor.go new file mode 100644 index 000000000..a7fa25c99 --- /dev/null +++ b/server/adka2a/v2/executor.go @@ -0,0 +1,459 @@ +// Copyright 2025 Google LLC +// +// Licensed under the Apache License, Version 2.0 (the "License"); +// you may not use this file except in compliance with the License. +// You may obtain a copy of the License at +// +// http://www.apache.org/licenses/LICENSE-2.0 +// +// Unless required by applicable law or agreed to in writing, software +// distributed under the License is distributed on an "AS IS" BASIS, +// WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. +// See the License for the specific language governing permissions and +// limitations under the License. + +package adka2a + +import ( + "context" + "errors" + "fmt" + "iter" + "slices" + + "github.com/a2aproject/a2a-go/v2/a2a" + "github.com/a2aproject/a2a-go/v2/a2asrv" + "github.com/a2aproject/a2a-go/v2/log" + + "google.golang.org/genai" + + "google.golang.org/adk/agent" + iremoteagent "google.golang.org/adk/internal/agent/remoteagent" + "google.golang.org/adk/plugin" + "google.golang.org/adk/runner" + "google.golang.org/adk/session" +) + +// BeforeExecuteCallback is the callback which will be called before an execution is started. +type BeforeExecuteCallback func(ctx context.Context, reqCtx *a2asrv.ExecutorContext) (context.Context, error) + +// AfterEventCallback is the callback which will be called after an ADK event is converted to an A2A event. +type AfterEventCallback func(ctx ExecutorContext, event *session.Event, processed *a2a.TaskArtifactUpdateEvent) error + +// AfterExecuteCallback is the callback which will be called after an execution resolved into a completed or failed task. +type AfterExecuteCallback func(ctx ExecutorContext, finalEvent *a2a.TaskStatusUpdateEvent, err error) error + +// A2APartConverter is a custom converter for converting A2A parts to GenAI parts. +// Implementations should generally remember to leverage adka2a.ToGenAiPart for default conversions +// nil returns are considered intentionally dropped parts. +type A2APartConverter func(ctx context.Context, a2aEvent a2a.Event, part *a2a.Part) (*genai.Part, error) + +// GenAIPartConverter is a custom converter for converting GenAI parts to A2A parts. +// Implementations should generally remember to leverage adka2a.ToA2APart for default conversions +// nil returns are considered intentionally dropped parts. +type GenAIPartConverter func(ctx context.Context, adkEvent *session.Event, part *genai.Part) (*a2a.Part, error) + +// A2AExecutionCleanupCallback is a callback which will be called after an execution or cancellation has completed or failed. +type A2AExecutionCleanupCallback func(ctx context.Context, reqCtx *a2asrv.ExecutorContext, subAgentCards []*a2a.AgentCard, result a2a.SendMessageResult, cause error) + +// OutputMode controls how artifacts are produced. +type OutputMode string + +const ( + // OutputArtifactPerRun produces a single artifact per [runner.Runner.Run]. + OutputArtifactPerRun OutputMode = "artifact-per-run" + // OutputArtifactPerEvent produces an artifact per non-partial [session.Event]. + // While agent is emitting events an artifact is build incrementally (parts are append to it). + // The next partial event replaces accumulated contents and seals the artifact, meaning + // the next event from this agent will create a new artifact. + OutputArtifactPerEvent OutputMode = "artifact-per-event" +) + +// Runner is an interface matching [runner.Runner] API. +// It exists to let users use custom runner implementations with A2A agent executor. +type Runner interface { + // Run runs the agent for the given user input, yielding events from agents. + Run(ctx context.Context, userID, sessionID string, msg *genai.Content, cfg agent.RunConfig) iter.Seq2[*session.Event, error] +} + +// RunnerProvider is a [Runner] factory function. The provided plugin must be installed in the returned [Runner] for +// callbacks taking [ExecutorContext] to work correctly. +type RunnerProvider func(ctx context.Context, reqCtx *a2asrv.ExecutorContext, plugin *plugin.Plugin) (RunnerConfig, Runner, error) + +// RunnerConfig is part of the runner configuration executor code depends on. +// Custom [RunnerProvider] needs to return it back to callers. +type RunnerConfig struct { + // AppName is the name of the application used in [session.Service] keys and A2A event metadata. + AppName string + // Agent is the root agent. + Agent agent.Agent + // SessionService is the session service to use. + SessionService session.Service +} + +// ExecutorConfig allows to configure Executor. +type ExecutorConfig struct { + // RunnerConfig is used for creating a default RunnerProvider. The field is ignored when RunnerProvider is set. + RunnerConfig runner.Config + // RunnerProvider is a function which allows to control how a runner is created. + // If not provided the default provider is used which calls [runner.New] with the RunnerConfig field. + RunnerProvider RunnerProvider + + // RunConfig is the configuration which will be passed to [runner.Runner.Run] during A2A Execute invocation. + RunConfig agent.RunConfig + + // BeforeExecuteCallback is the callback which will be called before an execution is started. + // It can be used to instrument a context or prevent the execution by returning an error. + BeforeExecuteCallback BeforeExecuteCallback + + // AfterEventCallback is the callback which will be called after an ADK event is successfully converted to an A2A event. + // This gives an opportunity to enrich the event with additional metadata or abort the execution by returning an error. + // The callback is not invoked for errors originating from ADK or event processing. Such errors are converted to + // TaskStatusUpdateEvent-s with TaskStateFailed state. If needed these can be intercepted using AfterExecuteCallback. + AfterEventCallback AfterEventCallback + + // AfterExecuteCallback is the callback which will be called after an execution resolved into a completed or failed task. + // This gives an opportunity to enrich the event with additional metadata or log it. + AfterExecuteCallback AfterExecuteCallback + + // A2APartConverter is a custom converter for converting A2A parts to GenAI parts. + // Implementations should generally remember to leverage [adka2a.ToGenAiPart] for default conversions + // nil returns are considered intentionally dropped parts. + A2APartConverter A2APartConverter + + // GenAIPartConverter is a custom converter for converting GenAI parts to A2A parts. + // Implementations should generally remember to leverage [adka2a.ToA2APart] for default conversions + // nil returns are considered intentionally dropped parts. + GenAIPartConverter GenAIPartConverter + + // OutputMode controls how artifacts are produced. Can be [OutputArtifactPerRun] or [OutputArtifactPerEvent]. + // Defaults to [OutputArtifactPerRun]. + OutputMode OutputMode + + // A2AExecutionCleanupCallback is a callback which will be called after an execution or cancellation has completed or failed. + // If not provided, the default behavior is to log the failure cause, if any. + A2AExecutionCleanupCallback A2AExecutionCleanupCallback +} + +var _ a2asrv.AgentExecutor = (*Executor)(nil) + +// Executor invokes an ADK agent and translates [session.Event]-s to [a2a.Event]-s according to the following rules: +// - If the input doesn't reference any a2a.Task, produce a Task with TaskStateSubmitted state. +// - Right before runner.Runner invocation, produce TaskStatusUpdateEvent with TaskStateWorking. +// - For every session.Event produce a TaskArtifactUpdateEvent{Append=true} with transformed parts. +// - After the last session.Event is processed produce an empty TaskArtifactUpdateEvent{Append=true} with LastChunk=true, +// if at least one artifact update was produced during the run. +// - If there was an LLMResponse with non-zero error code, produce a TaskStatusUpdateEvent with TaskStateFailed. +// Else if there was an LLMResponse with long-running tool invocation, produce a TaskStatusUpdateEvent with TaskStateInputRequired. +// Else produce a TaskStatusUpdateEvent with TaskStateCompleted. +type Executor struct { + config ExecutorConfig +} + +// NewExecutor creates an initialized [Executor] instance. +func NewExecutor(config ExecutorConfig) *Executor { + if config.RunnerProvider == nil { + config.RunnerProvider = newDefaultRunnerProvider(config.RunnerConfig) + } + return &Executor{config: config} +} + +func (e *Executor) Execute(ctx context.Context, execCtx *a2asrv.ExecutorContext) iter.Seq2[a2a.Event, error] { + return func(yield func(a2a.Event, error) bool) { + msg := execCtx.Message + if msg == nil { + yield(nil, fmt.Errorf("message not provided")) + return + } + content, err := toGenAIContent(ctx, msg, e.config.A2APartConverter) + if err != nil { + yield(nil, fmt.Errorf("a2a message conversion failed: %w", err)) + return + } + + executorPlugin, err := newExecutorPlugin() + if err != nil { + yield(nil, fmt.Errorf("failed to create a2a-executor plugin: %w", err)) + return + } + + cfg, r, err := e.config.RunnerProvider(ctx, execCtx, executorPlugin.plugin) + if err != nil { + yield(nil, fmt.Errorf("failed to create a runner: %w", err)) + return + } + if e.config.BeforeExecuteCallback != nil { + ctx, err = e.config.BeforeExecuteCallback(ctx, execCtx) + if err != nil { + yield(nil, fmt.Errorf("before execute: %w", err)) + return + } + if ctx == nil { + yield(nil, fmt.Errorf("before execute: callback returned nil context")) + return + } + } + + if event, err := HandleInputRequired(execCtx, content); event != nil || err != nil { + if err != nil { + yield(nil, err) + } else { + yield(event, nil) + } + return + } + + if execCtx.StoredTask == nil { + event := a2a.NewSubmittedTask(execCtx, msg) + if !yield(event, nil) { + return + } + } + + invocationMeta := toInvocationMeta(ctx, cfg, execCtx) + + err = e.prepareSession(ctx, cfg, invocationMeta) + if err != nil { + statusEvent := toTaskFailedUpdateEvent(execCtx, err, invocationMeta.eventMeta) + execCtx := newExecutorContext(ctx, invocationMeta, executorPlugin, content) + e.writeFinalTaskStatus(execCtx, yield, nil, statusEvent, err) + return + } + + event := a2a.NewStatusUpdateEvent(execCtx, a2a.TaskStateWorking, nil) + event.Metadata = invocationMeta.eventMeta + if !yield(event, nil) { + return + } + + var artifactTransform eventToArtifactTransform + if e.config.OutputMode == OutputArtifactPerEvent { + artifactTransform = newArtifactMaker(execCtx) + } else { + artifactTransform = newLegacyArtifactMaker(execCtx) + } + + processor := newEventProcessor(execCtx, invocationMeta, e.config.GenAIPartConverter, artifactTransform) + executorContext := newExecutorContext(ctx, invocationMeta, executorPlugin, content) + e.process(executorContext, r, processor, yield) + } +} + +func (e *Executor) Cancel(ctx context.Context, execCtx *a2asrv.ExecutorContext) iter.Seq2[a2a.Event, error] { + return func(yield func(a2a.Event, error) bool) { + event := a2a.NewStatusUpdateEvent(execCtx, a2a.TaskStateCanceled, nil) + yield(event, nil) + } +} + +func (e *Executor) Cleanup(ctx context.Context, execCtx *a2asrv.ExecutorContext, result a2a.SendMessageResult, cause error) { + cfg, err := e.createRunnerConfig(ctx, execCtx) + if err != nil { + log.Error(ctx, "failed to create runner config", err) + + if e.config.A2AExecutionCleanupCallback != nil { + e.config.A2AExecutionCleanupCallback(ctx, execCtx, []*a2a.AgentCard{}, result, cause) + } + return + } + + remoteSubagents := findRemoteSubagents(cfg.Agent) + + // If task was in input-required and got successfully cancelled - run the cleanup logic + if execCtx.StoredTask != nil && execCtx.StoredTask.Status.State == a2a.TaskStateInputRequired { + if task, ok := result.(*a2a.Task); ok && task.Status.State == a2a.TaskStateCanceled && execCtx.Message == nil { + if err := e.cancelChildInputRequiredTasks(ctx, execCtx, execCtx.StoredTask.Status, cfg, remoteSubagents); err != nil { + log.Warn(ctx, "failed to cancel subagent tasks waiting for input", "cause", err) + } + } + } + + if e.config.A2AExecutionCleanupCallback != nil { + subAgentCards := make([]*a2a.AgentCard, len(remoteSubagents)) + for i, subagent := range remoteSubagents { + subAgentCards[i] = subagent.config.AgentCard + } + e.config.A2AExecutionCleanupCallback(ctx, execCtx, subAgentCards, result, cause) + } else if cause != nil { + if execCtx.Message != nil { + log.Warn(ctx, "execution failed", "error", cause) + } else { + log.Warn(ctx, "cancellation failed", "error", cause) + } + } +} + +func (e *Executor) cancelChildInputRequiredTasks(ctx context.Context, reqCtx *a2asrv.ExecutorContext, status a2a.TaskStatus, cfg RunnerConfig, subagents []remoteAgent) error { + if len(subagents) == 0 { + return nil + } + + meta := toInvocationMeta(ctx, cfg, reqCtx) + getSessionResponse, err := cfg.SessionService.Get(ctx, &session.GetRequest{ + AppName: cfg.AppName, + UserID: meta.userID, + SessionID: meta.sessionID, + }) + if err != nil { + return fmt.Errorf("failed to get a session: %w", err) + } + + tasksToCancel, err := getSubagentTasksToCancel(ctx, status, getSessionResponse.Session) + if err != nil { + return fmt.Errorf("subtask search failed: %w", err) + } + if len(tasksToCancel) == 0 { + return nil + } + + var failures []error + clientCache := map[string]iremoteagent.A2AClient{} + for _, task := range tasksToCancel { // TODO(yarolegovich): run in parallel (how to limit?) + remoteSubagentIdx := slices.IndexFunc(subagents, func(a remoteAgent) bool { return a.agent.Name() == task.agentName }) + if remoteSubagentIdx < 0 { + continue + } + remoteSubagent := subagents[remoteSubagentIdx] + client, ok := clientCache[task.agentName] + if !ok { + _, newClient, err := iremoteagent.CreateA2AClient(ctx, remoteSubagent.config) + if err != nil { + failures = append(failures, fmt.Errorf("failed to create A2A client: %w", err)) + continue + } + clientCache[task.agentName] = newClient + client = newClient + } + _, err = client.CancelTask(ctx, &a2a.CancelTaskRequest{ID: task.taskID}) + if err != nil { + failures = append(failures, fmt.Errorf("failed to cancel task: %w", err)) + continue + } + } + for _, client := range clientCache { + if err := client.Destroy(); err != nil { + failures = append(failures, fmt.Errorf("client destroy failed: %w", err)) + } + } + return errors.Join(failures...) +} + +// Processing failures should be delivered as Task failed events. An error is returned from this method if an event write fails. +func (e *Executor) process(ctx ExecutorContext, r Runner, processor *eventProcessor, yield func(a2a.Event, error) bool) { + meta := processor.meta + for adkEvent, adkErr := range r.Run(ctx, meta.userID, meta.sessionID, ctx.UserContent(), e.config.RunConfig) { + if adkErr != nil { + event := processor.makeTaskFailedEvent(ctx, fmt.Errorf("agent run failed: %w", adkErr), nil) + e.writeFinalTaskStatus(ctx, yield, processor.makeFinalArtifactUpdate(), event, adkErr) + return + } + + a2aEvent, pErr := processor.process(ctx, adkEvent) + if pErr == nil && a2aEvent != nil && e.config.AfterEventCallback != nil { + pErr = e.config.AfterEventCallback(ctx, adkEvent, a2aEvent) + } + + if pErr != nil { + event := processor.makeTaskFailedEvent(ctx, fmt.Errorf("processor failed: %w", pErr), adkEvent) + e.writeFinalTaskStatus(ctx, yield, processor.makeFinalArtifactUpdate(), event, pErr) + return + } + + if a2aEvent != nil { + if !yield(a2aEvent, nil) { + return + } + } + } + + finalStatus := processor.makeFinalStatusUpdate() + e.writeFinalTaskStatus(ctx, yield, processor.makeFinalArtifactUpdate(), finalStatus, nil) +} + +func (e *Executor) writeFinalTaskStatus( + ctx ExecutorContext, + yield func(a2a.Event, error) bool, + partialReset *a2a.TaskArtifactUpdateEvent, + status *a2a.TaskStatusUpdateEvent, + err error, +) { + if e.config.AfterExecuteCallback != nil { + if err = e.config.AfterExecuteCallback(ctx, status, err); err != nil { + yield(nil, fmt.Errorf("after execute: %w", err)) + return + } + } + if partialReset != nil { + if !yield(partialReset, nil) { + return + } + } + yield(status, nil) +} + +func (e *Executor) prepareSession(ctx context.Context, cfg RunnerConfig, meta invocationMeta) error { + service := cfg.SessionService + + _, err := service.Get(ctx, &session.GetRequest{ + AppName: cfg.AppName, + UserID: meta.userID, + SessionID: meta.sessionID, + }) + if err == nil { + return nil + } + + _, err = service.Create(ctx, &session.CreateRequest{ + AppName: cfg.AppName, + UserID: meta.userID, + SessionID: meta.sessionID, + State: make(map[string]any), + }) + if err != nil { + return fmt.Errorf("failed to create a session: %w", err) + } + return nil +} + +func (e *Executor) createRunnerConfig(ctx context.Context, reqCtx *a2asrv.ExecutorContext) (RunnerConfig, error) { + executorPlugin, err := newExecutorPlugin() + if err != nil { + return RunnerConfig{}, fmt.Errorf("failed to create a2a-plugin: %w", err) + } + cfg, _, err := e.config.RunnerProvider(ctx, reqCtx, executorPlugin.plugin) + if err != nil { + return RunnerConfig{}, fmt.Errorf("runner provider failed: %w", err) + } + return cfg, nil +} + +func newDefaultRunnerProvider(baseConfig runner.Config) RunnerProvider { + return func(ctx context.Context, reqCtx *a2asrv.ExecutorContext, plugin *plugin.Plugin) (RunnerConfig, Runner, error) { + if baseConfig.Agent == nil { + return RunnerConfig{}, nil, fmt.Errorf("runner.Config.Agent is not provided") + } + if baseConfig.SessionService == nil { + return RunnerConfig{}, nil, fmt.Errorf("runner.Config.SessionService is not provided") + } + + cfg := baseConfig + cfg.PluginConfig.Plugins = append(slices.Clone(cfg.PluginConfig.Plugins), plugin) + r, err := runner.New(cfg) + if err != nil { + return RunnerConfig{}, nil, err + } + return toInternalRunnerConfig(cfg), &defaultRunner{runner: r}, nil + } +} + +type defaultRunner struct { + runner *runner.Runner +} + +func (r *defaultRunner) Run(ctx context.Context, userID, sessionID string, msg *genai.Content, cfg agent.RunConfig) iter.Seq2[*session.Event, error] { + return r.runner.Run(ctx, userID, sessionID, msg, cfg) +} + +func toInternalRunnerConfig(cfg runner.Config) RunnerConfig { + return RunnerConfig{Agent: cfg.Agent, AppName: cfg.AppName, SessionService: cfg.SessionService} +} diff --git a/server/adka2a/executor_context.go b/server/adka2a/v2/executor_context.go similarity index 92% rename from server/adka2a/executor_context.go rename to server/adka2a/v2/executor_context.go index 33a54b207..54345547d 100644 --- a/server/adka2a/executor_context.go +++ b/server/adka2a/v2/executor_context.go @@ -18,7 +18,7 @@ import ( "context" "iter" - "github.com/a2aproject/a2a-go/a2asrv" + "github.com/a2aproject/a2a-go/v2/a2asrv" "google.golang.org/genai" "google.golang.org/adk/session" @@ -42,8 +42,8 @@ type ExecutorContext interface { Events() session.Events // UserContent is a converted A2A message which is passed to runner.Run. UserContent() *genai.Content - // RequestContext containts information about the original A2A Request, the current task and related tasks. - RequestContext() *a2asrv.RequestContext + // RequestContext contains information about the original A2A Request, the current task and related tasks. + RequestContext() *a2asrv.ExecutorContext } type executorContext struct { @@ -90,7 +90,7 @@ func (ec *executorContext) Events() session.Events { return session.Events() } -func (ec *executorContext) RequestContext() *a2asrv.RequestContext { +func (ec *executorContext) RequestContext() *a2asrv.ExecutorContext { return ec.meta.reqCtx } diff --git a/server/adka2a/executor_plugin.go b/server/adka2a/v2/executor_plugin.go similarity index 100% rename from server/adka2a/executor_plugin.go rename to server/adka2a/v2/executor_plugin.go diff --git a/server/adka2a/executor_test.go b/server/adka2a/v2/executor_test.go similarity index 70% rename from server/adka2a/executor_test.go rename to server/adka2a/v2/executor_test.go index 60525bb92..4457c398a 100644 --- a/server/adka2a/executor_test.go +++ b/server/adka2a/v2/executor_test.go @@ -22,10 +22,9 @@ import ( "testing" "time" - "github.com/a2aproject/a2a-go/a2a" - "github.com/a2aproject/a2a-go/a2aclient" - "github.com/a2aproject/a2a-go/a2asrv" - "github.com/a2aproject/a2a-go/a2asrv/eventqueue" + "github.com/a2aproject/a2a-go/v2/a2a" + "github.com/a2aproject/a2a-go/v2/a2aclient" + "github.com/a2aproject/a2a-go/v2/a2asrv" "github.com/google/go-cmp/cmp" "github.com/google/go-cmp/cmp/cmpopts" "google.golang.org/genai" @@ -38,19 +37,9 @@ import ( "google.golang.org/adk/session" ) -type testQueue struct { - eventqueue.Queue - events []a2a.Event - writeErr *eventIndex -} +// No testQueue needed anymore. -func (q *testQueue) Write(_ context.Context, e a2a.Event) error { - if q.writeErr != nil && q.writeErr.i == len(q.events) { - return fmt.Errorf("queue write failed") - } - q.events = append(q.events, e) - return nil -} +// testQueue removed. type testSessionService struct { session.Service @@ -82,26 +71,16 @@ func newEventReplayAgent(events []*session.Event, failWith error) (agent.Agent, }) } -func newInMemoryQueue(t *testing.T) eventqueue.Queue { - t.Helper() - qm := eventqueue.NewInMemoryManager() - q, err := qm.GetOrCreate(t.Context(), "test") - if err != nil { - t.Fatalf("qm.GetOrCreate() error = %v", err) - } - return q -} - type eventIndex struct{ i int } func TestExecutor_Execute(t *testing.T) { task := &a2a.Task{ID: a2a.NewTaskID(), ContextID: a2a.NewContextID()} - hiMsg := a2a.NewMessage(a2a.MessageRoleUser, a2a.TextPart{Text: "hi"}) - hiMsgForTask := a2a.NewMessageForTask(a2a.MessageRoleUser, task, a2a.TextPart{Text: "hi"}) + hiMsg := a2a.NewMessage(a2a.MessageRoleUser, a2a.NewTextPart("hi")) + hiMsgForTask := a2a.NewMessageForTask(a2a.MessageRoleUser, task, a2a.NewTextPart("hi")) testCases := []struct { name string - request *a2a.MessageSendParams + request *a2a.SendMessageRequest events []*session.Event wantEvents []a2a.Event createSessionFails bool @@ -111,30 +90,37 @@ func TestExecutor_Execute(t *testing.T) { }{ { name: "no message", - request: &a2a.MessageSendParams{}, + request: &a2a.SendMessageRequest{}, wantErr: true, }, { name: "malformed data", - request: &a2a.MessageSendParams{Message: a2a.NewMessageForTask(a2a.MessageRoleUser, task, a2a.FilePart{ - File: a2a.FileBytes{Bytes: "(*_*)"}, // malformed base64 - })}, + request: &a2a.SendMessageRequest{ + Message: a2a.NewMessageForTask(a2a.MessageRoleUser, task, func() *a2a.Part { + p := a2a.NewDataPart(make(chan int)) + p.SetMeta(a2aDataPartMetaTypeKey, a2aDataPartTypeFunctionCall) + return p + }()), + }, wantErr: true, }, { name: "session setup fails", - request: &a2a.MessageSendParams{Message: hiMsgForTask}, + request: &a2a.SendMessageRequest{Message: hiMsgForTask}, createSessionFails: true, wantEvents: []a2a.Event{ newFinalStatusUpdate( task, a2a.TaskStateFailed, - a2a.NewMessageForTask(a2a.MessageRoleAgent, task, a2a.TextPart{Text: "failed to create a session: session creation failed"}), + a2a.NewMessageForTask(a2a.MessageRoleAgent, task, &a2a.Part{ + Content: a2a.Text("failed to create a session: session creation failed"), + Metadata: map[string]any{metadataIsErrMessageKey: true}, + }), ), }, }, { name: "success for a new task", - request: &a2a.MessageSendParams{Message: hiMsg}, + request: &a2a.SendMessageRequest{Message: hiMsg}, events: []*session.Event{ {LLMResponse: modelResponseFromParts(genai.NewPartFromText("Hello"))}, {LLMResponse: modelResponseFromParts(genai.NewPartFromText(", world!"))}, @@ -142,34 +128,34 @@ func TestExecutor_Execute(t *testing.T) { wantEvents: []a2a.Event{ a2a.NewSubmittedTask(task, hiMsg), a2a.NewStatusUpdateEvent(task, a2a.TaskStateWorking, nil), - a2a.NewArtifactEvent(task, a2a.TextPart{Text: "Hello"}), - a2a.NewArtifactUpdateEvent(task, a2a.NewArtifactID(), a2a.TextPart{Text: ", world!"}), + a2a.NewArtifactEvent(task, a2a.NewTextPart("Hello")), + a2a.NewArtifactUpdateEvent(task, a2a.NewArtifactID(), a2a.NewTextPart(", world!")), newFinalStatusUpdate(task, a2a.TaskStateCompleted, nil), }, }, { name: "success for existing task", - request: &a2a.MessageSendParams{Message: hiMsgForTask}, + request: &a2a.SendMessageRequest{Message: hiMsgForTask}, events: []*session.Event{ {LLMResponse: modelResponseFromParts(genai.NewPartFromText("Hello"))}, {LLMResponse: modelResponseFromParts(genai.NewPartFromText(", world!"))}, }, wantEvents: []a2a.Event{ a2a.NewStatusUpdateEvent(task, a2a.TaskStateWorking, nil), - a2a.NewArtifactEvent(task, a2a.TextPart{Text: "Hello"}), - a2a.NewArtifactUpdateEvent(task, a2a.NewArtifactID(), a2a.TextPart{Text: ", world!"}), + a2a.NewArtifactEvent(task, a2a.NewTextPart("Hello")), + a2a.NewArtifactUpdateEvent(task, a2a.NewArtifactID(), a2a.NewTextPart(", world!")), newFinalStatusUpdate(task, a2a.TaskStateCompleted, nil), }, }, { name: "queue write fails", - request: &a2a.MessageSendParams{Message: hiMsgForTask}, + request: &a2a.SendMessageRequest{Message: hiMsgForTask}, queueWriteFails: &eventIndex{0}, wantErr: true, }, { name: "llm fails", - request: &a2a.MessageSendParams{Message: hiMsgForTask}, + request: &a2a.SendMessageRequest{Message: hiMsgForTask}, events: []*session.Event{ {LLMResponse: modelResponseFromParts(genai.NewPartFromText("Hello"))}, {LLMResponse: model.LLMResponse{ErrorCode: "418", ErrorMessage: "I'm a teapot"}}, @@ -177,8 +163,8 @@ func TestExecutor_Execute(t *testing.T) { }, wantEvents: []a2a.Event{ a2a.NewStatusUpdateEvent(task, a2a.TaskStateWorking, nil), - a2a.NewArtifactEvent(task, a2a.TextPart{Text: "Hello"}), - a2a.NewArtifactUpdateEvent(task, a2a.NewArtifactID(), a2a.TextPart{Text: ", world!"}), + a2a.NewArtifactEvent(task, a2a.NewTextPart("Hello")), + a2a.NewArtifactUpdateEvent(task, a2a.NewArtifactID(), a2a.NewTextPart(", world!")), toTaskFailedUpdateEvent( task, errorFromResponse(&model.LLMResponse{ErrorCode: "418", ErrorMessage: "I'm a teapot"}), map[string]any{ToA2AMetaKey("error_code"): "418"}, @@ -187,23 +173,26 @@ func TestExecutor_Execute(t *testing.T) { }, { name: "agent run fails", - request: &a2a.MessageSendParams{Message: hiMsgForTask}, + request: &a2a.SendMessageRequest{Message: hiMsgForTask}, events: []*session.Event{ {LLMResponse: modelResponseFromParts(genai.NewPartFromText("Hello"))}, }, agentRunFails: fmt.Errorf("OOF"), wantEvents: []a2a.Event{ a2a.NewStatusUpdateEvent(task, a2a.TaskStateWorking, nil), - a2a.NewArtifactEvent(task, a2a.TextPart{Text: "Hello"}), + a2a.NewArtifactEvent(task, a2a.NewTextPart("Hello")), newFinalStatusUpdate( task, a2a.TaskStateFailed, - a2a.NewMessageForTask(a2a.MessageRoleAgent, task, a2a.TextPart{Text: "agent run failed: OOF"}), + a2a.NewMessageForTask(a2a.MessageRoleAgent, task, &a2a.Part{ + Content: a2a.Text("agent run failed: OOF"), + Metadata: map[string]any{metadataIsErrMessageKey: true}, + }), ), }, }, { name: "agent run and queue write fail", - request: &a2a.MessageSendParams{Message: hiMsgForTask}, + request: &a2a.SendMessageRequest{Message: hiMsgForTask}, events: []*session.Event{ {LLMResponse: modelResponseFromParts(genai.NewPartFromText("Hello"))}, }, @@ -212,13 +201,13 @@ func TestExecutor_Execute(t *testing.T) { wantErr: true, wantEvents: []a2a.Event{ a2a.NewStatusUpdateEvent(task, a2a.TaskStateWorking, nil), - a2a.NewArtifactEvent(task, a2a.TextPart{Text: "Hello"}), + a2a.NewArtifactEvent(task, a2a.NewTextPart("Hello")), }, }, } for _, tc := range testCases { - ignoreFields := []cmp.Option{ + ignoreOptions := []cmp.Option{ cmpopts.IgnoreFields(a2a.Message{}, "ID"), cmpopts.IgnoreFields(a2a.Artifact{}, "ID"), cmpopts.IgnoreFields(a2a.TaskStatus{}, "Timestamp"), @@ -234,23 +223,35 @@ func TestExecutor_Execute(t *testing.T) { sessionService := &testSessionService{Service: session.InMemoryService(), createErr: tc.createSessionFails} runnerConfig := runner.Config{AppName: agent.Name(), Agent: agent, SessionService: sessionService} executor := NewExecutor(ExecutorConfig{RunnerConfig: runnerConfig}) - queue := &testQueue{Queue: newInMemoryQueue(t), writeErr: tc.queueWriteFails} - reqCtx := &a2asrv.RequestContext{TaskID: task.ID, ContextID: task.ContextID, Message: tc.request.Message} + var gotEvents []a2a.Event + reqCtx := &a2asrv.ExecutorContext{TaskID: task.ID, ContextID: task.ContextID, Message: tc.request.Message} if tc.request.Message != nil && tc.request.Message.TaskID == task.ID { reqCtx.StoredTask = task } - err = executor.Execute(t.Context(), reqCtx, queue) - if err != nil && !tc.wantErr { - t.Fatalf("executor.Execute() error = %v, want nil", err) - } - if err == nil && tc.wantErr { - t.Fatalf("executor.Execute() produced %d events, want error", len(queue.events)) + var executeErr error + for event, err := range executor.Execute(t.Context(), reqCtx) { + if err != nil { + executeErr = err + break + } + if tc.queueWriteFails != nil && tc.queueWriteFails.i == len(gotEvents) { + executeErr = fmt.Errorf("queue write failed") + break + } + gotEvents = append(gotEvents, event) } - if tc.wantEvents != nil { - if diff := cmp.Diff(tc.wantEvents, queue.events, ignoreFields...); diff != "" { - t.Fatalf("executor.Execute() wrong events (+got,-want):\ngot = %v\nwant = %v\ndiff = %s", queue.events, tc.wantEvents, diff) + if executeErr != nil { + if !tc.wantErr { + t.Fatalf("executor.Execute() error = %v, want nil", executeErr) } + return + } + if tc.wantErr { + t.Fatalf("executor.Execute() error = nil, want error") + } + if diff := cmp.Diff(tc.wantEvents, gotEvents, ignoreOptions...); diff != "" { + t.Errorf("executor.Execute() produced wrong events (-want +got):\n%s", diff) } }) } @@ -259,19 +260,21 @@ func TestExecutor_Execute(t *testing.T) { func TestExecutor_Cancel(t *testing.T) { task := &a2a.Task{ID: a2a.NewTaskID(), ContextID: a2a.NewContextID()} executor := NewExecutor(ExecutorConfig{}) - reqCtx := &a2asrv.RequestContext{TaskID: task.ID, ContextID: task.ContextID} + reqCtx := &a2asrv.ExecutorContext{TaskID: task.ID, ContextID: task.ContextID} - queue := &testQueue{Queue: newInMemoryQueue(t)} + var gotEvents []a2a.Event reqCtx.StoredTask = task - err := executor.Cancel(t.Context(), reqCtx, queue) - if err != nil { - t.Fatalf("executor.Cancel() error = %v, want nil", err) + for event, err := range executor.Cancel(t.Context(), reqCtx) { + if err != nil { + t.Fatalf("executor.Cancel() error = %v, want nil", err) + } + gotEvents = append(gotEvents, event) } - if len(queue.events) != 1 { - t.Fatalf("executor.Cancel() produced %d events, want 1", queue.events) + if len(gotEvents) != 1 { + t.Fatalf("executor.Cancel() produced %d events, want 1", len(gotEvents)) } - event := queue.events[0].(*a2a.TaskStatusUpdateEvent) + event := gotEvents[0].(*a2a.TaskStatusUpdateEvent) if event.Status.State != a2a.TaskStateCanceled { t.Fatalf("executor.Cancel() = %v, want a single TaskStateCanceled update", event) } @@ -286,20 +289,24 @@ func TestExecutor_SessionReuse(t *testing.T) { sessionService := session.InMemoryService() task := &a2a.Task{ID: a2a.NewTaskID(), ContextID: a2a.NewContextID()} - req := &a2a.MessageSendParams{Message: a2a.NewMessageForTask(a2a.MessageRoleUser, task)} - reqCtx := &a2asrv.RequestContext{TaskID: task.ID, ContextID: task.ContextID, Message: req.Message} + req := &a2a.SendMessageRequest{Message: a2a.NewMessageForTask(a2a.MessageRoleUser, task)} + reqCtx := &a2asrv.ExecutorContext{TaskID: task.ID, ContextID: task.ContextID, Message: req.Message} runnerConfig := runner.Config{AppName: agent.Name(), Agent: agent, SessionService: sessionService} config := ExecutorConfig{RunnerConfig: runnerConfig} executor := NewExecutor(config) - queue := newInMemoryQueue(t) + var gotEvents []a2a.Event - err = executor.Execute(ctx, reqCtx, queue) - if err != nil { - t.Fatalf("executor.Execute() error = %v, want nil", err) + for event, err := range executor.Execute(ctx, reqCtx) { + if err != nil { + t.Fatalf("executor.Execute() error = %v, want nil", err) + } + gotEvents = append(gotEvents, event) } - err = executor.Execute(ctx, reqCtx, queue) - if err != nil { - t.Fatalf("executor.Execute() error = %v, want nil", err) + for event, err := range executor.Execute(ctx, reqCtx) { + if err != nil { + t.Fatalf("executor.Execute() error = %v, want nil", err) + } + gotEvents = append(gotEvents, event) } meta := toInvocationMeta(ctx, toInternalRunnerConfig(config.RunnerConfig), reqCtx) @@ -321,7 +328,7 @@ func TestExecutor_SessionReuse(t *testing.T) { func TestExecutor_Callbacks(t *testing.T) { type contextKeyType struct{} task := &a2a.Task{ID: a2a.NewTaskID(), ContextID: a2a.NewContextID()} - hiMsg := a2a.NewMessageForTask(a2a.MessageRoleUser, task, a2a.TextPart{Text: "hi"}) + hiMsg := a2a.NewMessageForTask(a2a.MessageRoleUser, task, a2a.NewTextPart("hi")) testCases := []struct { name string @@ -335,24 +342,24 @@ func TestExecutor_Callbacks(t *testing.T) { }{ { name: "abort execution", - beforeExecution: func(ctx context.Context, reqCtx *a2asrv.RequestContext) (context.Context, error) { + beforeExecution: func(ctx context.Context, reqCtx *a2asrv.ExecutorContext) (context.Context, error) { return nil, fmt.Errorf("aborted") }, wantErr: fmt.Errorf("aborted"), }, { name: "instrument context", - beforeExecution: func(ctx context.Context, reqCtx *a2asrv.RequestContext) (context.Context, error) { + beforeExecution: func(ctx context.Context, reqCtx *a2asrv.ExecutorContext) (context.Context, error) { return context.WithValue(ctx, contextKeyType{}, "bar"), nil }, afterExecution: func(ctx ExecutorContext, finalUpdate *a2a.TaskStatusUpdateEvent, err error) error { text, _ := ctx.Value(contextKeyType{}).(string) - finalUpdate.Status.Message = a2a.NewMessage(a2a.MessageRoleAgent, a2a.TextPart{Text: text}) + finalUpdate.Status.Message = a2a.NewMessage(a2a.MessageRoleAgent, a2a.NewTextPart(text)) return nil }, wantEvents: []a2a.Event{ a2a.NewStatusUpdateEvent(task, a2a.TaskStateWorking, nil), - newFinalStatusUpdate(task, a2a.TaskStateCompleted, a2a.NewMessage(a2a.MessageRoleAgent, a2a.TextPart{Text: "bar"})), + newFinalStatusUpdate(task, a2a.TaskStateCompleted, a2a.NewMessage(a2a.MessageRoleAgent, a2a.NewTextPart("bar"))), }, }, { @@ -364,12 +371,12 @@ func TestExecutor_Callbacks(t *testing.T) { return fmt.Errorf("fail!") }, afterExecution: func(ctx ExecutorContext, finalUpdate *a2a.TaskStatusUpdateEvent, err error) error { - finalUpdate.Status.Message = a2a.NewMessage(a2a.MessageRoleAgent, a2a.TextPart{Text: "bar"}) + finalUpdate.Status.Message = a2a.NewMessage(a2a.MessageRoleAgent, a2a.NewTextPart("bar")) return nil }, wantEvents: []a2a.Event{ a2a.NewStatusUpdateEvent(task, a2a.TaskStateWorking, nil), - newFinalStatusUpdate(task, a2a.TaskStateFailed, a2a.NewMessage(a2a.MessageRoleAgent, a2a.TextPart{Text: "bar"})), + newFinalStatusUpdate(task, a2a.TaskStateFailed, a2a.NewMessage(a2a.MessageRoleAgent, a2a.NewTextPart("bar"))), }, }, { @@ -380,10 +387,10 @@ func TestExecutor_Callbacks(t *testing.T) { for range ctx.ReadonlyState().All() { eventCount++ } - finalUpdate.Status.Message = a2a.NewMessage(a2a.MessageRoleAgent, a2a.TextPart{Text: fmt.Sprintf("%d events", eventCount)}) + finalUpdate.Status.Message = a2a.NewMessage(a2a.MessageRoleAgent, a2a.NewTextPart(fmt.Sprintf("%d events", eventCount))) return nil }, - wantEvents: []a2a.Event{newFinalStatusUpdate(task, a2a.TaskStateFailed, a2a.NewMessage(a2a.MessageRoleAgent, a2a.TextPart{Text: "0 events"}))}, + wantEvents: []a2a.Event{newFinalStatusUpdate(task, a2a.TaskStateFailed, a2a.NewMessage(a2a.MessageRoleAgent, a2a.NewTextPart("0 events")))}, }, { name: "enrich event", @@ -392,13 +399,13 @@ func TestExecutor_Callbacks(t *testing.T) { {LLMResponse: modelResponseFromParts(genai.NewPartFromText(", world!"))}, }, afterEvent: func(ctx ExecutorContext, event *session.Event, processed *a2a.TaskArtifactUpdateEvent) error { - processed.Artifact.Parts = append(processed.Artifact.Parts, a2a.TextPart{Text: " (enriched)"}) + processed.Artifact.Parts = append(processed.Artifact.Parts, a2a.NewTextPart(" (enriched)")) return nil }, wantEvents: []a2a.Event{ a2a.NewStatusUpdateEvent(task, a2a.TaskStateWorking, nil), - a2a.NewArtifactEvent(task, a2a.TextPart{Text: "Hello"}, a2a.TextPart{Text: " (enriched)"}), - a2a.NewArtifactUpdateEvent(task, a2a.NewArtifactID(), a2a.TextPart{Text: ", world!"}, a2a.TextPart{Text: " (enriched)"}), + a2a.NewArtifactEvent(task, a2a.NewTextPart("Hello"), a2a.NewTextPart(" (enriched)")), + a2a.NewArtifactUpdateEvent(task, a2a.NewArtifactID(), a2a.NewTextPart(", world!"), a2a.NewTextPart(" (enriched)")), newFinalStatusUpdate(task, a2a.TaskStateCompleted, nil), }, }, @@ -410,15 +417,15 @@ func TestExecutor_Callbacks(t *testing.T) { }, afterExecution: func(ctx ExecutorContext, finalUpdate *a2a.TaskStatusUpdateEvent, err error) error { eventCount := ctx.Events().Len() - finalUpdate.Status.Message = a2a.NewMessage(a2a.MessageRoleAgent, a2a.TextPart{Text: fmt.Sprintf("event count = %d", eventCount)}) + finalUpdate.Status.Message = a2a.NewMessage(a2a.MessageRoleAgent, a2a.NewTextPart(fmt.Sprintf("event count = %d", eventCount))) return nil }, wantEvents: []a2a.Event{ a2a.NewStatusUpdateEvent(task, a2a.TaskStateWorking, nil), - a2a.NewArtifactEvent(task, a2a.TextPart{Text: "Hello"}), - a2a.NewArtifactUpdateEvent(task, a2a.NewArtifactID(), a2a.TextPart{Text: ", world!"}), + a2a.NewArtifactEvent(task, a2a.NewTextPart("Hello")), + a2a.NewArtifactUpdateEvent(task, a2a.NewArtifactID(), a2a.NewTextPart(", world!")), newFinalStatusUpdate(task, a2a.TaskStateCompleted, - a2a.NewMessage(a2a.MessageRoleAgent, a2a.TextPart{Text: "event count = 3"}), + a2a.NewMessage(a2a.MessageRoleAgent, a2a.NewTextPart("event count = 3")), ), }, }, @@ -439,7 +446,7 @@ func TestExecutor_Callbacks(t *testing.T) { } for _, tc := range testCases { - ignoreFields := []cmp.Option{ + ignoreOptions := []cmp.Option{ cmpopts.IgnoreFields(a2a.Message{}, "ID"), cmpopts.IgnoreFields(a2a.Artifact{}, "ID"), cmpopts.IgnoreFields(a2a.TaskStatus{}, "Timestamp"), @@ -460,19 +467,29 @@ func TestExecutor_Callbacks(t *testing.T) { AfterEventCallback: tc.afterEvent, AfterExecuteCallback: tc.afterExecution, }) - queue := &testQueue{Queue: newInMemoryQueue(t)} - reqCtx := &a2asrv.RequestContext{TaskID: task.ID, ContextID: task.ContextID, Message: hiMsg, StoredTask: task} - - err = executor.Execute(t.Context(), reqCtx, queue) - if err != nil && tc.wantErr == nil { - t.Fatalf("executor.Execute() error = %v, want nil", err) + var gotEvents []a2a.Event + reqCtx := &a2asrv.ExecutorContext{TaskID: task.ID, ContextID: task.ContextID, Message: hiMsg, StoredTask: task} + + var executeErr error + for event, err := range executor.Execute(t.Context(), reqCtx) { + if err != nil { + executeErr = err + break + } + gotEvents = append(gotEvents, event) } - if err == nil && tc.wantErr != nil { + if executeErr != nil { + if tc.wantErr == nil { + t.Fatalf("executor.Execute() error = %v, want nil", executeErr) + } + return + } + if tc.wantErr != nil { t.Fatalf("executor.Execute() error = nil, want %v", tc.wantErr) } if tc.wantEvents != nil { - if diff := cmp.Diff(tc.wantEvents, queue.events, ignoreFields...); diff != "" { - t.Fatalf("executor.Execute() wrong events (+got,-want):\ngot = %v\nwant = %v\ndiff = %s", queue.events, tc.wantEvents, diff) + if diff := cmp.Diff(tc.wantEvents, gotEvents, ignoreOptions...); diff != "" { + t.Errorf("executor.Execute() produced wrong events (-want +got):\n%s", diff) } } }) @@ -514,9 +531,9 @@ func TestExecutor_Cancel_AfterEvent(t *testing.T) { defer server.Close() card := &a2a.AgentCard{ - Name: "test", - URL: server.URL, - PreferredTransport: a2a.TransportProtocolJSONRPC, + SupportedInterfaces: []*a2a.AgentInterface{ + a2a.NewAgentInterface(server.URL, a2a.TransportProtocolJSONRPC), + }, } client, err := a2aclient.NewFromCard(t.Context(), card) @@ -525,11 +542,9 @@ func TestExecutor_Cancel_AfterEvent(t *testing.T) { } msgId := a2a.NewMessageID() - blocking := false - - result, sendErr := client.SendMessage(t.Context(), &a2a.MessageSendParams{ - Message: &a2a.Message{ID: string(msgId), Parts: []a2a.Part{a2a.TextPart{Text: "TEST"}}, Role: a2a.MessageRoleUser}, - Config: &a2a.MessageSendConfig{Blocking: &blocking}, + result, sendErr := client.SendMessage(t.Context(), &a2a.SendMessageRequest{ + Message: &a2a.Message{ID: string(msgId), Parts: a2a.ContentParts{a2a.NewTextPart("TEST")}, Role: a2a.MessageRoleUser}, + Config: &a2a.SendMessageConfig{ReturnImmediately: true}, }) if sendErr != nil { @@ -538,7 +553,7 @@ func TestExecutor_Cancel_AfterEvent(t *testing.T) { taskID := result.TaskInfo().TaskID - task, err := client.CancelTask(t.Context(), &a2a.TaskIDParams{ID: taskID}) + task, err := client.CancelTask(t.Context(), &a2a.CancelTaskRequest{ID: taskID}) if err != nil { t.Fatalf("client.CancelTask() error = %v, want nil", err) } @@ -558,7 +573,7 @@ func TestExecutor_Cancel_AfterEvent(t *testing.T) { func TestExecutor_Converters(t *testing.T) { task := &a2a.Task{ID: a2a.NewTaskID(), ContextID: a2a.NewContextID()} - hiMsg := a2a.NewMessageForTask(a2a.MessageRoleUser, task, a2a.TextPart{Text: "hi"}) + hiMsg := a2a.NewMessageForTask(a2a.MessageRoleUser, task, a2a.NewTextPart("hi")) t.Run("A2APartConverter", func(t *testing.T) { t.Run("modify input", func(t *testing.T) { @@ -578,17 +593,20 @@ func TestExecutor_Converters(t *testing.T) { executor := NewExecutor(ExecutorConfig{ RunnerConfig: runner.Config{AppName: agent.Name(), Agent: agent, SessionService: session.InMemoryService()}, - A2APartConverter: func(ctx context.Context, evt a2a.Event, part a2a.Part) (*genai.Part, error) { - if p, ok := part.(a2a.TextPart); ok && p.Text == "hi" { + A2APartConverter: func(ctx context.Context, evt a2a.Event, part *a2a.Part) (*genai.Part, error) { + if part.Text() == "hi" { return genai.NewPartFromText("HELLO"), nil } return nil, nil }, }) - reqCtx := &a2asrv.RequestContext{TaskID: task.ID, ContextID: task.ContextID, Message: hiMsg, StoredTask: task} - if err := executor.Execute(t.Context(), reqCtx, newInMemoryQueue(t)); err != nil { - t.Fatalf("executor.Execute() error = %v", err) + reqCtx := &a2asrv.ExecutorContext{TaskID: task.ID, ContextID: task.ContextID, Message: hiMsg, StoredTask: task} + for event, err := range executor.Execute(t.Context(), reqCtx) { + if err != nil { + t.Fatalf("executor.Execute() error = %v", err) + } + _ = event // ignored } if receivedText != "HELLO" { @@ -611,14 +629,17 @@ func TestExecutor_Converters(t *testing.T) { executor := NewExecutor(ExecutorConfig{ RunnerConfig: runner.Config{AppName: agent.Name(), Agent: agent, SessionService: session.InMemoryService()}, - A2APartConverter: func(ctx context.Context, evt a2a.Event, part a2a.Part) (*genai.Part, error) { + A2APartConverter: func(ctx context.Context, evt a2a.Event, part *a2a.Part) (*genai.Part, error) { return nil, nil }, }) - reqCtx := &a2asrv.RequestContext{TaskID: task.ID, ContextID: task.ContextID, Message: hiMsg, StoredTask: task} - if err := executor.Execute(t.Context(), reqCtx, newInMemoryQueue(t)); err != nil { - t.Fatalf("executor.Execute() error = %v", err) + reqCtx := &a2asrv.ExecutorContext{TaskID: task.ID, ContextID: task.ContextID, Message: hiMsg, StoredTask: task} + for event, err := range executor.Execute(t.Context(), reqCtx) { + if err != nil { + t.Fatalf("executor.Execute() error = %v", err) + } + _ = event // ignored } if receivedParts != 0 { @@ -642,32 +663,35 @@ func TestExecutor_Converters(t *testing.T) { executor := NewExecutor(ExecutorConfig{ RunnerConfig: runner.Config{AppName: agent.Name(), Agent: agent, SessionService: session.InMemoryService()}, - GenAIPartConverter: func(ctx context.Context, evt *session.Event, part *genai.Part) (a2a.Part, error) { + GenAIPartConverter: func(ctx context.Context, evt *session.Event, part *genai.Part) (*a2a.Part, error) { if part.Text == "world" { - return a2a.TextPart{Text: "WORLD"}, nil + return a2a.NewTextPart("WORLD"), nil } - return a2a.TextPart{Text: part.Text}, nil + return a2a.NewTextPart(part.Text), nil }, }) - queue := &testQueue{Queue: newInMemoryQueue(t)} - reqCtx := &a2asrv.RequestContext{TaskID: task.ID, ContextID: task.ContextID, Message: hiMsg, StoredTask: task} - if err := executor.Execute(t.Context(), reqCtx, queue); err != nil { - t.Fatalf("executor.Execute() error = %v", err) + var gotEvents []a2a.Event + reqCtx := &a2asrv.ExecutorContext{TaskID: task.ID, ContextID: task.ContextID, Message: hiMsg, StoredTask: task} + for event, err := range executor.Execute(t.Context(), reqCtx) { + if err != nil { + t.Fatalf("executor.Execute() error = %v", err) + } + gotEvents = append(gotEvents, event) } found := false - for _, e := range queue.events { + for _, e := range gotEvents { if ae, ok := e.(*a2a.TaskArtifactUpdateEvent); ok { for _, p := range ae.Artifact.Parts { - if tp, ok := p.(a2a.TextPart); ok && tp.Text == "WORLD" { + if p.Text() == "WORLD" { found = true } } } } if !found { - t.Errorf("did not find 'WORLD' in events: %v", queue.events) + t.Errorf("did not find 'WORLD' in events: %v", gotEvents) } }) @@ -679,18 +703,21 @@ func TestExecutor_Converters(t *testing.T) { executor := NewExecutor(ExecutorConfig{ RunnerConfig: runner.Config{AppName: agent.Name(), Agent: agent, SessionService: session.InMemoryService()}, - GenAIPartConverter: func(ctx context.Context, evt *session.Event, part *genai.Part) (a2a.Part, error) { + GenAIPartConverter: func(ctx context.Context, evt *session.Event, part *genai.Part) (*a2a.Part, error) { return nil, nil }, }) - queue := &testQueue{Queue: newInMemoryQueue(t)} - reqCtx := &a2asrv.RequestContext{TaskID: task.ID, ContextID: task.ContextID, Message: hiMsg, StoredTask: task} - if err := executor.Execute(t.Context(), reqCtx, queue); err != nil { - t.Fatalf("executor.Execute() error = %v", err) + var gotEvents []a2a.Event + reqCtx := &a2asrv.ExecutorContext{TaskID: task.ID, ContextID: task.ContextID, Message: hiMsg, StoredTask: task} + for event, err := range executor.Execute(t.Context(), reqCtx) { + if err != nil { + t.Fatalf("executor.Execute() error = %v", err) + } + gotEvents = append(gotEvents, event) } - for _, e := range queue.events { + for _, e := range gotEvents { if ae, ok := e.(*a2a.TaskArtifactUpdateEvent); ok { if len(ae.Artifact.Parts) != 0 { t.Errorf("found %d parts but expected 0", len(ae.Artifact.Parts)) @@ -703,7 +730,7 @@ func TestExecutor_Converters(t *testing.T) { func TestExecutor_OutputArtifactPerEvent(t *testing.T) { task := &a2a.Task{ID: a2a.NewTaskID(), ContextID: a2a.NewContextID()} - hiMsg := a2a.NewMessageForTask(a2a.MessageRoleUser, task, a2a.TextPart{Text: "hi"}) + hiMsg := a2a.NewMessageForTask(a2a.MessageRoleUser, task, a2a.NewTextPart("hi")) testCases := []struct { name string @@ -721,15 +748,15 @@ func TestExecutor_OutputArtifactPerEvent(t *testing.T) { wantEvents: []a2a.Event{ a2a.NewStatusUpdateEvent(task, a2a.TaskStateWorking, nil), &a2a.TaskArtifactUpdateEvent{ - Artifact: &a2a.Artifact{Parts: a2a.ContentParts{a2a.TextPart{Text: "Hello, "}}}, + Artifact: &a2a.Artifact{Parts: a2a.ContentParts{a2a.NewTextPart("Hello, ")}}, Append: false, LastChunk: false, }, &a2a.TaskArtifactUpdateEvent{ - Artifact: &a2a.Artifact{Parts: a2a.ContentParts{a2a.TextPart{Text: "world!"}}}, + Artifact: &a2a.Artifact{Parts: a2a.ContentParts{a2a.NewTextPart("world!")}}, Append: true, LastChunk: false, }, &a2a.TaskArtifactUpdateEvent{ - Artifact: &a2a.Artifact{Parts: a2a.ContentParts{a2a.TextPart{Text: "Hello, world!"}}}, + Artifact: &a2a.Artifact{Parts: a2a.ContentParts{a2a.NewTextPart("Hello, world!")}}, Append: false, LastChunk: true, }, newFinalStatusUpdate(task, a2a.TaskStateCompleted, nil), @@ -746,19 +773,19 @@ func TestExecutor_OutputArtifactPerEvent(t *testing.T) { wantEvents: []a2a.Event{ a2a.NewStatusUpdateEvent(task, a2a.TaskStateWorking, nil), &a2a.TaskArtifactUpdateEvent{ - Artifact: &a2a.Artifact{Parts: a2a.ContentParts{a2a.TextPart{Text: "Agent1: H"}}}, + Artifact: &a2a.Artifact{Parts: a2a.ContentParts{a2a.NewTextPart("Agent1: H")}}, Append: false, LastChunk: false, }, &a2a.TaskArtifactUpdateEvent{ - Artifact: &a2a.Artifact{Parts: a2a.ContentParts{a2a.TextPart{Text: "Agent2: W"}}}, + Artifact: &a2a.Artifact{Parts: a2a.ContentParts{a2a.NewTextPart("Agent2: W")}}, Append: false, LastChunk: false, }, &a2a.TaskArtifactUpdateEvent{ - Artifact: &a2a.Artifact{Parts: a2a.ContentParts{a2a.TextPart{Text: "Agent1: Hello"}}}, + Artifact: &a2a.Artifact{Parts: a2a.ContentParts{a2a.NewTextPart("Agent1: Hello")}}, Append: false, LastChunk: true, }, &a2a.TaskArtifactUpdateEvent{ - Artifact: &a2a.Artifact{Parts: a2a.ContentParts{a2a.TextPart{Text: "Agent2: World"}}}, + Artifact: &a2a.Artifact{Parts: a2a.ContentParts{a2a.NewTextPart("Agent2: World")}}, Append: false, LastChunk: true, }, newFinalStatusUpdate(task, a2a.TaskStateCompleted, nil), @@ -782,16 +809,19 @@ func TestExecutor_OutputArtifactPerEvent(t *testing.T) { a2a.NewStatusUpdateEvent(task, a2a.TaskStateWorking, nil), &a2a.TaskArtifactUpdateEvent{ Artifact: &a2a.Artifact{ - Parts: a2a.ContentParts{a2a.TextPart{ - Text: "Thinking...", - Metadata: map[string]any{ToA2AMetaKey("thought"): true}, - }}, + Parts: a2a.ContentParts{ + func() *a2a.Part { + p := a2a.NewTextPart("Thinking...") + p.SetMeta(ToA2AMetaKey("thought"), true) + return p + }(), + }, }, Append: false, LastChunk: false, }, &a2a.TaskArtifactUpdateEvent{ Artifact: &a2a.Artifact{ - Parts: a2a.ContentParts{a2a.TextPart{Text: "Done"}}, + Parts: a2a.ContentParts{a2a.NewTextPart("Done")}, }, Append: false, LastChunk: true, }, @@ -818,11 +848,12 @@ func TestExecutor_OutputArtifactPerEvent(t *testing.T) { &a2a.TaskArtifactUpdateEvent{ Artifact: &a2a.Artifact{ Parts: a2a.ContentParts{ - a2a.TextPart{ - Text: "Thought part", - Metadata: map[string]any{ToA2AMetaKey("thought"): true}, - }, - a2a.TextPart{Text: "Text part"}, + func() *a2a.Part { + p := a2a.NewTextPart("Thought part") + p.SetMeta(ToA2AMetaKey("thought"), true) + return p + }(), + a2a.NewTextPart("Text part"), }, }, Append: false, LastChunk: true, @@ -841,19 +872,19 @@ func TestExecutor_OutputArtifactPerEvent(t *testing.T) { wantEvents: []a2a.Event{ a2a.NewStatusUpdateEvent(task, a2a.TaskStateWorking, nil), &a2a.TaskArtifactUpdateEvent{ - Artifact: &a2a.Artifact{Parts: a2a.ContentParts{a2a.TextPart{Text: "1.a"}}}, + Artifact: &a2a.Artifact{Parts: a2a.ContentParts{a2a.NewTextPart("1.a")}}, Append: false, LastChunk: false, }, &a2a.TaskArtifactUpdateEvent{ - Artifact: &a2a.Artifact{Parts: a2a.ContentParts{a2a.TextPart{Text: "1.final"}}}, + Artifact: &a2a.Artifact{Parts: a2a.ContentParts{a2a.NewTextPart("1.final")}}, Append: false, LastChunk: true, }, &a2a.TaskArtifactUpdateEvent{ - Artifact: &a2a.Artifact{Parts: a2a.ContentParts{a2a.TextPart{Text: "2.a"}}}, + Artifact: &a2a.Artifact{Parts: a2a.ContentParts{a2a.NewTextPart("2.a")}}, Append: false, LastChunk: false, }, &a2a.TaskArtifactUpdateEvent{ - Artifact: &a2a.Artifact{Parts: a2a.ContentParts{a2a.TextPart{Text: "2.final"}}}, + Artifact: &a2a.Artifact{Parts: a2a.ContentParts{a2a.NewTextPart("2.final")}}, Append: false, LastChunk: true, }, newFinalStatusUpdate(task, a2a.TaskStateCompleted, nil), @@ -869,11 +900,14 @@ func TestExecutor_OutputArtifactPerEvent(t *testing.T) { OutputMode: OutputArtifactPerEvent, }) - queue := &testQueue{Queue: newInMemoryQueue(t)} - reqCtx := &a2asrv.RequestContext{TaskID: task.ID, ContextID: task.ContextID, Message: hiMsg, StoredTask: task} + var gotEvents []a2a.Event + reqCtx := &a2asrv.ExecutorContext{TaskID: task.ID, ContextID: task.ContextID, Message: hiMsg, StoredTask: task} - if err := executor.Execute(t.Context(), reqCtx, queue); err != nil { - t.Fatalf("executor.Execute() error = %v", err) + for event, err := range executor.Execute(t.Context(), reqCtx) { + if err != nil { + t.Fatalf("executor.Execute() error = %v", err) + } + gotEvents = append(gotEvents, event) } ignoreOptions := []cmp.Option{ @@ -883,15 +917,15 @@ func TestExecutor_OutputArtifactPerEvent(t *testing.T) { cmpopts.IgnoreFields(a2a.TaskStatusUpdateEvent{}, "Metadata", "TaskID", "ContextID"), cmpopts.IgnoreFields(a2a.TaskArtifactUpdateEvent{}, "Metadata", "TaskID", "ContextID"), } - if len(queue.events) != len(tc.wantEvents) { - t.Fatalf("got %d events, want %d", len(queue.events), len(tc.wantEvents)) + if len(gotEvents) != len(tc.wantEvents) { + t.Fatalf("got %d events, want %d", len(gotEvents), len(tc.wantEvents)) } authorToID, lastFinishedID := make(map[string]a2a.ArtifactID), make(map[string]a2a.ArtifactID) artifactCount := 0 - for i := range queue.events { - got := queue.events[i] + for i := range gotEvents { + got := gotEvents[i] want := tc.wantEvents[i] if diff := cmp.Diff(want, got, ignoreOptions...); diff != "" { @@ -942,8 +976,8 @@ func TestExecutor_RunnerProvider(t *testing.T) { wantText := "Hello" ctx := t.Context() task := &a2a.Task{ID: a2a.NewTaskID(), ContextID: a2a.NewContextID()} - hiMsg := a2a.NewMessageForTask(a2a.MessageRoleUser, task, a2a.TextPart{Text: "hi"}) - reqCtx := &a2asrv.RequestContext{TaskID: task.ID, ContextID: task.ContextID, Message: hiMsg, StoredTask: task} + hiMsg := a2a.NewMessageForTask(a2a.MessageRoleUser, task, a2a.NewTextPart("hi")) + reqCtx := &a2asrv.ExecutorContext{TaskID: task.ID, ContextID: task.ContextID, Message: hiMsg, StoredTask: task} runnerConfig := runner.Config{ AppName: "test", @@ -952,7 +986,7 @@ func TestExecutor_RunnerProvider(t *testing.T) { } executor := NewExecutor(ExecutorConfig{ RunnerConfig: runnerConfig, - RunnerProvider: func(pCtx context.Context, pReqCtx *a2asrv.RequestContext, plugin *plugin.Plugin) (RunnerConfig, Runner, error) { + RunnerProvider: func(pCtx context.Context, pReqCtx *a2asrv.ExecutorContext, plugin *plugin.Plugin) (RunnerConfig, Runner, error) { return toInternalRunnerConfig(runnerConfig), &testRunner{ runFunc: func(ctx context.Context, userID, sessionID string, msg *genai.Content, cfg agent.RunConfig) iter.Seq2[*session.Event, error] { return func(yield func(*session.Event, error) bool) { @@ -963,16 +997,22 @@ func TestExecutor_RunnerProvider(t *testing.T) { }, }) - queue := &testQueue{Queue: newInMemoryQueue(t)} - if err := executor.Execute(ctx, reqCtx, queue); err != nil { - t.Fatalf("executor.Execute() error = %v", err) + var events []a2a.Event + for event, err := range executor.Execute(ctx, reqCtx) { + if err != nil { + t.Fatalf("executor.Execute() error = %v", err) + } + events = append(events, event) + } + if len(events) < 2 { + t.Fatalf("executor.Execute() produced %d events, want at least 2", len(events)) } - ta, ok := queue.events[1].(*a2a.TaskArtifactUpdateEvent) + ta, ok := events[1].(*a2a.TaskArtifactUpdateEvent) if !ok { - t.Fatalf("queue.events[1] = %T, want a2a.TaskArtifactUpdateEvent", queue.events[1]) + t.Fatalf("queue.events[1] = %T, want a2a.TaskArtifactUpdateEvent", events[1]) } - if tp, ok := ta.Artifact.Parts[0].(a2a.TextPart); !ok || tp.Text != wantText { - t.Fatalf("ta.Artifact.Parts[0] = %v, want text part with text = %q", tp, wantText) + if ta.Artifact.Parts[0].Text() != wantText { + t.Fatalf("ta.Artifact.Parts[0] = %v, want text part with text = %q", ta.Artifact.Parts[0], wantText) } } diff --git a/server/adka2a/input_required.go b/server/adka2a/v2/input_required.go similarity index 78% rename from server/adka2a/input_required.go rename to server/adka2a/v2/input_required.go index d6d57e78e..a2a6671dd 100644 --- a/server/adka2a/input_required.go +++ b/server/adka2a/v2/input_required.go @@ -19,23 +19,24 @@ import ( "fmt" "slices" - "github.com/a2aproject/a2a-go/a2a" - "github.com/a2aproject/a2a-go/a2asrv" - "github.com/a2aproject/a2a-go/log" + "github.com/a2aproject/a2a-go/v2/a2a" + "github.com/a2aproject/a2a-go/v2/a2asrv" + "github.com/a2aproject/a2a-go/v2/log" "google.golang.org/genai" + "google.golang.org/adk/internal/utils" "google.golang.org/adk/session" ) type inputRequiredProcessor struct { - reqCtx *a2asrv.RequestContext + reqCtx *a2asrv.ExecutorContext event *a2a.TaskStatusUpdateEvent partConverter GenAIPartConverter // handles possible duplication in partial and non-partial events addedParts []*genai.Part } -func newInputRequiredProcessor(reqCtx *a2asrv.RequestContext, partConverter GenAIPartConverter) *inputRequiredProcessor { +func newInputRequiredProcessor(reqCtx *a2asrv.ExecutorContext, partConverter GenAIPartConverter) *inputRequiredProcessor { return &inputRequiredProcessor{reqCtx: reqCtx, partConverter: partConverter} } @@ -88,7 +89,6 @@ func (p *inputRequiredProcessor) process(ctx context.Context, event *session.Eve } else { msg := a2a.NewMessage(a2a.MessageRoleAgent, a2aParts...) ev := a2a.NewStatusUpdateEvent(p.reqCtx, a2a.TaskStateInputRequired, msg) - ev.Final = true p.event = ev } } @@ -105,11 +105,11 @@ func (p *inputRequiredProcessor) process(ctx context.Context, event *session.Eve return &modifiedEvent, nil } -func (p *inputRequiredProcessor) convertParts(ctx context.Context, event *session.Event, parts []*genai.Part, longRunningCallIDs []string) ([]a2a.Part, error) { +func (p *inputRequiredProcessor) convertParts(ctx context.Context, event *session.Event, parts []*genai.Part, longRunningCallIDs []string) ([]*a2a.Part, error) { if p.partConverter == nil { return ToA2AParts(parts, longRunningCallIDs) } - converted := make([]a2a.Part, 0, len(parts)) + converted := make([]*a2a.Part, 0, len(parts)) for _, part := range parts { cp, err := p.partConverter(ctx, event, part) if err != nil { @@ -135,9 +135,9 @@ func (p *inputRequiredProcessor) isLongRunningResponse(event *session.Event, par return false } for _, part := range p.event.Status.Message.Parts { - if dp, ok := part.(a2a.DataPart); ok { - if typeVal, ok := dp.Metadata[a2aDataPartMetaTypeKey]; ok && typeVal == a2aDataPartTypeFunctionCall { - if callID, ok := dp.Data["id"].(string); ok && callID == id { + if data, ok := part.Data().(map[string]any); ok { + if typeVal, ok := part.Meta()[a2aDataPartMetaTypeKey]; ok && typeVal == a2aDataPartTypeFunctionCall { + if callID, ok := data["id"].(string); ok && callID == id { return true } } @@ -146,10 +146,10 @@ func (p *inputRequiredProcessor) isLongRunningResponse(event *session.Event, par return false } -// handleInputRequired checks if the input message contains responses to all function calls +// HandleInputRequired checks if the input message contains responses to all function calls // that happened during the previous invocation and were recorded in the Task input-required state message. // If a non-nil event is returned the invoking code needs to use the event as the result of the execution -func handleInputRequired(reqCtx *a2asrv.RequestContext, content *genai.Content) (*a2a.TaskStatusUpdateEvent, error) { +func HandleInputRequired(reqCtx *a2asrv.ExecutorContext, content *genai.Content) (*a2a.TaskStatusUpdateEvent, error) { if reqCtx.StoredTask == nil { return nil, nil } @@ -174,7 +174,6 @@ func handleInputRequired(reqCtx *a2asrv.RequestContext, content *genai.Content) parts := makeInputMissingErrorMessage(statusMsg.Parts, statusPart.FunctionCall.ID) msg := a2a.NewMessageForTask(a2a.MessageRoleAgent, reqCtx.StoredTask, parts...) event := a2a.NewStatusUpdateEvent(reqCtx, a2a.TaskStateInputRequired, msg) - event.Final = true return event, nil } } @@ -190,28 +189,36 @@ func getSubagentTasksToCancel(ctx context.Context, status a2a.TaskStatus, sessio return nil, nil } - foundCalls := 0 + foundCalls, foundTaskIDs := map[string]bool{}, 0 var tasksToCancel []subagentTask events := session.Events() for i := events.Len() - 1; i >= 0; i-- { event := events.At(i) - if event.Content != nil { - if slices.ContainsFunc(event.Content.Parts, func(p *genai.Part) bool { - return p.FunctionCall != nil && slices.Contains(pendingCallIDs, p.FunctionCall.ID) - }) { - foundCalls++ - if taskID, _ := GetA2ATaskInfo(event); taskID != "" { - tasksToCancel = append(tasksToCancel, subagentTask{agentName: event.Author, taskID: a2a.TaskID(taskID)}) - } + for _, call := range utils.FunctionCalls(event.Content) { + if !slices.Contains(pendingCallIDs, call.ID) { + continue + } + if foundCalls[call.ID] { + log.Warn(ctx, "duplicate function call id", "id", call.ID) + continue + } + foundCalls[call.ID] = true + if taskID, _ := GetA2ATaskInfo(event); taskID != "" { + tasksToCancel = append(tasksToCancel, subagentTask{agentName: event.Author, taskID: a2a.TaskID(taskID)}) + foundTaskIDs++ } } - if foundCalls == len(pendingCallIDs) { + if len(foundCalls) == len(pendingCallIDs) { break } } - if foundCalls < len(pendingCallIDs) { - log.Warn(ctx, "could not find all function calls from status message", "found", foundCalls, "total", len(pendingCallIDs)) + if len(foundCalls) < len(pendingCallIDs) { + log.Warn(ctx, "could not find all function calls from status message", "found", len(foundCalls), "total", len(pendingCallIDs)) + } + + if foundTaskIDs < len(foundCalls) { + log.Warn(ctx, "not all function calls had a taskID", "found_calls", len(foundCalls), "found_ids", foundTaskIDs) } return tasksToCancel, nil @@ -235,12 +242,11 @@ func getPendingLongRunningCallIDs(status a2a.TaskStatus) ([]string, error) { return callIDs, nil } -func makeInputMissingErrorMessage(inputRequiredParts []a2a.Part, callID string) []a2a.Part { - errPart := a2a.TextPart{ - Text: fmt.Sprintf("no input provided for function call ID %q", callID), - Metadata: map[string]any{"validation_error": true}, - } - var preservedParts []a2a.Part +func makeInputMissingErrorMessage(inputRequiredParts []*a2a.Part, callID string) []*a2a.Part { + errPart := a2a.NewTextPart(fmt.Sprintf("no input provided for function call ID %q", callID)) + errPart.SetMeta("validation_error", true) + + var preservedParts []*a2a.Part for _, p := range inputRequiredParts { if meta := p.Meta(); meta != nil { if v, ok := meta["validation_error"].(bool); ok && v { diff --git a/server/adka2a/metadata.go b/server/adka2a/v2/metadata.go similarity index 92% rename from server/adka2a/metadata.go rename to server/adka2a/v2/metadata.go index fb29367be..d7159cb03 100644 --- a/server/adka2a/metadata.go +++ b/server/adka2a/v2/metadata.go @@ -18,8 +18,8 @@ import ( "context" "maps" - "github.com/a2aproject/a2a-go/a2a" - "github.com/a2aproject/a2a-go/a2asrv" + "github.com/a2aproject/a2a-go/v2/a2a" + "github.com/a2aproject/a2a-go/v2/a2asrv" "google.golang.org/adk/internal/converters" "google.golang.org/adk/session" @@ -37,14 +37,15 @@ var ( metadataUsageKey = ToA2AMetaKey("usage_metadata") metadataCustomMetaKey = ToA2AMetaKey("custom_metadata") metadataPartialKey = ToA2AMetaKey("partial") + metadataIsErrMessageKey = ToA2AMetaKey("is_error_message") ) -// ToA2AMetaKey adds a prefix used to differentiage ADK-related values stored in Metadata an A2A event. +// ToA2AMetaKey adds a prefix used to differentiate ADK-related values stored in Metadata an A2A event. func ToA2AMetaKey(key string) string { return "adk_" + key } -// ToADKMetaKey adds a prefix used to differentiage A2A-related values stored in custom metadata of an ADK session event. +// ToADKMetaKey adds a prefix used to differentiate A2A-related values stored in custom metadata of an ADK session event. func ToADKMetaKey(key string) string { return "a2a:" + key } @@ -53,17 +54,17 @@ type invocationMeta struct { userID string sessionID string agentName string - reqCtx *a2asrv.RequestContext + reqCtx *a2asrv.ExecutorContext eventMeta map[string]any } -func toInvocationMeta(ctx context.Context, config RunnerConfig, reqCtx *a2asrv.RequestContext) invocationMeta { +func toInvocationMeta(ctx context.Context, config RunnerConfig, reqCtx *a2asrv.ExecutorContext) invocationMeta { userID, sessionID := "A2A_USER_"+reqCtx.ContextID, reqCtx.ContextID // a2a sdk attaches authn info to the call context, use it when provided if callCtx, ok := a2asrv.CallContextFrom(ctx); ok { - if callCtx.User != nil && callCtx.User.Name() != "" { - userID = callCtx.User.Name() + if callCtx.User != nil && callCtx.User.Name != "" { + userID = callCtx.User.Name } } diff --git a/server/adka2a/metadata_test.go b/server/adka2a/v2/metadata_test.go similarity index 99% rename from server/adka2a/metadata_test.go rename to server/adka2a/v2/metadata_test.go index 03c92d79b..d0e2c1da8 100644 --- a/server/adka2a/metadata_test.go +++ b/server/adka2a/v2/metadata_test.go @@ -17,7 +17,7 @@ package adka2a import ( "testing" - "github.com/a2aproject/a2a-go/a2a" + "github.com/a2aproject/a2a-go/v2/a2a" "github.com/google/go-cmp/cmp" "google.golang.org/genai" diff --git a/server/adka2a/parts.go b/server/adka2a/v2/parts.go similarity index 65% rename from server/adka2a/parts.go rename to server/adka2a/v2/parts.go index c693bceb3..ea22ab5d6 100644 --- a/server/adka2a/parts.go +++ b/server/adka2a/v2/parts.go @@ -17,13 +17,12 @@ package adka2a import ( "bytes" "context" - "encoding/base64" "encoding/json" "fmt" "maps" "slices" - "github.com/a2aproject/a2a-go/a2a" + "github.com/a2aproject/a2a-go/v2/a2a" "google.golang.org/genai" "google.golang.org/adk/internal/converters" @@ -41,7 +40,7 @@ const ( a2aDataPartTypeCodeExecutableCode = "executable_code" ) -// IsPartial takes metadata of an A2A object (eg. a2a.Part, a2a.Artifact) and returs true if +// IsPartial takes metadata of an A2A object (eg. a2a.Part, a2a.Artifact) and returns true if // it was marked as partial based on the ADK partial flag set on the original ADK object. func IsPartial(meta map[string]any) bool { if meta == nil { @@ -51,7 +50,7 @@ func IsPartial(meta map[string]any) bool { return isPartial } -// IsPartialFlagSet takes metadata of an A2A object (eg. a2a.Part, a2a.Artifact) and returs true if +// IsPartialFlagSet takes metadata of an A2A object (eg. a2a.Part, a2a.Artifact) and returns true if // the ADK partial flag was set on it. func IsPartialFlagSet(meta map[string]any) bool { if meta == nil { @@ -75,7 +74,7 @@ func validateDataPartJSON(d *genai.Part) ([]byte, bool) { // ToA2APart converts the provided genai part to A2A equivalent. Long running tool IDs are used for attaching metadata to // the relevant data parts. -func ToA2APart(part *genai.Part, longRunningToolIDs []string) (a2a.Part, error) { +func ToA2APart(part *genai.Part, longRunningToolIDs []string) (*a2a.Part, error) { parts, err := ToA2AParts([]*genai.Part{part}, longRunningToolIDs) if err != nil { return nil, err @@ -85,13 +84,13 @@ func ToA2APart(part *genai.Part, longRunningToolIDs []string) (a2a.Part, error) // ToA2AParts converts the provided genai parts to A2A equivalents. Long running tool IDs are used for attaching metadata to // the relevant data parts. -func ToA2AParts(parts []*genai.Part, longRunningToolIDs []string) ([]a2a.Part, error) { - result := make([]a2a.Part, len(parts)) +func ToA2AParts(parts []*genai.Part, longRunningToolIDs []string) ([]*a2a.Part, error) { + result := make([]*a2a.Part, len(parts)) for i, part := range parts { if part.Text != "" { - r := a2a.TextPart{Text: part.Text} + r := a2a.NewTextPart(part.Text) if part.Thought { - r.Metadata = map[string]any{ToA2AMetaKey("thought"): true} + r.SetMeta(ToA2AMetaKey("thought"), true) } result[i] = r } else if jsonBytes, ok := validateDataPartJSON(part); ok { @@ -99,7 +98,7 @@ func ToA2AParts(parts []*genai.Part, longRunningToolIDs []string) ([]a2a.Part, e if err := json.Unmarshal(jsonBytes, &data); err != nil { return nil, err } - result[i] = a2a.DataPart{Data: data} + result[i] = a2a.NewDataPart(data) } else if part.InlineData != nil || part.FileData != nil { r, err := toA2AFilePart(part) if err != nil { @@ -117,84 +116,50 @@ func ToA2AParts(parts []*genai.Part, longRunningToolIDs []string) ([]a2a.Part, e return result, nil } -func updatePartsMetadata(parts []a2a.Part, update map[string]any) { - for i, part := range parts { - var meta map[string]any - switch p := part.(type) { - case a2a.TextPart: - if p.Metadata == nil { - p.Metadata = make(map[string]any) - parts[i] = p - } - meta = p.Metadata - case a2a.FilePart: - if p.Metadata == nil { - p.Metadata = make(map[string]any) - parts[i] = p - } - meta = p.Metadata - case a2a.DataPart: - if p.Metadata == nil { - p.Metadata = make(map[string]any) - parts[i] = p - } - meta = p.Metadata - default: - // TODO: log unknown part type warning (should never happen) - continue +func updatePartsMetadata(parts []*a2a.Part, update map[string]any) { + for _, part := range parts { + if part.Metadata == nil { + part.Metadata = make(map[string]any) } - maps.Copy(meta, update) + maps.Copy(part.Metadata, update) } } -func toA2AFilePart(v *genai.Part) (a2a.FilePart, error) { +func toA2AFilePart(v *genai.Part) (*a2a.Part, error) { if v == nil || (v.FileData == nil && v.InlineData == nil) { - return a2a.FilePart{}, fmt.Errorf("not a file part: %v", v) + return nil, fmt.Errorf("not a file part: %v", v) } if v.FileData != nil { - return a2a.FilePart{ - File: a2a.FileURI{ - FileMeta: a2a.FileMeta{ - Name: v.FileData.DisplayName, - MimeType: v.FileData.MIMEType, - }, - URI: v.FileData.FileURI, - }, - }, nil + r := a2a.NewFileURLPart(a2a.URL(v.FileData.FileURI), v.FileData.MIMEType) + r.Filename = v.FileData.DisplayName + return r, nil } - part := a2a.FilePart{ - File: a2a.FileBytes{ - FileMeta: a2a.FileMeta{ - Name: v.InlineData.DisplayName, - MimeType: v.InlineData.MIMEType, - }, - Bytes: base64.StdEncoding.EncodeToString(v.InlineData.Data), - }, - } + r := a2a.NewRawPart(v.InlineData.Data) + r.MediaType = v.InlineData.MIMEType + r.Filename = v.InlineData.DisplayName if v.VideoMetadata != nil { data, err := converters.ToMapStructure(v.VideoMetadata) if err != nil { - return a2a.FilePart{}, err + return nil, err } - part.Metadata = map[string]any{"video_metadata": data} + r.SetMeta("video_metadata", data) } - return part, nil + return r, nil } -func toA2ADataPart(part *genai.Part, longRunningToolIDs []string) (a2a.Part, error) { +func toA2ADataPart(part *genai.Part, longRunningToolIDs []string) (*a2a.Part, error) { if part.CodeExecutionResult != nil { data, err := converters.ToMapStructure(part.CodeExecutionResult) if err != nil { return nil, err } - return a2a.DataPart{ - Data: data, - Metadata: map[string]any{a2aDataPartMetaTypeKey: a2aDataPartTypeCodeExecResult}, - }, nil + r := a2a.NewDataPart(data) + r.SetMeta(a2aDataPartMetaTypeKey, a2aDataPartTypeCodeExecResult) + return r, nil } if part.FunctionResponse != nil { @@ -202,10 +167,9 @@ func toA2ADataPart(part *genai.Part, longRunningToolIDs []string) (a2a.Part, err if err != nil { return nil, err } - return a2a.DataPart{ - Data: data, - Metadata: map[string]any{a2aDataPartMetaTypeKey: a2aDataPartTypeFunctionResponse}, - }, nil + r := a2a.NewDataPart(data) + r.SetMeta(a2aDataPartMetaTypeKey, a2aDataPartTypeFunctionResponse) + return r, nil } if part.ExecutableCode != nil { @@ -213,10 +177,9 @@ func toA2ADataPart(part *genai.Part, longRunningToolIDs []string) (a2a.Part, err if err != nil { return nil, err } - return a2a.DataPart{ - Data: data, - Metadata: map[string]any{a2aDataPartMetaTypeKey: a2aDataPartTypeCodeExecutableCode}, - }, nil + r := a2a.NewDataPart(data) + r.SetMeta(a2aDataPartMetaTypeKey, a2aDataPartTypeCodeExecutableCode) + return r, nil } if part.FunctionCall != nil { @@ -224,20 +187,17 @@ func toA2ADataPart(part *genai.Part, longRunningToolIDs []string) (a2a.Part, err if err != nil { return nil, err } - return a2a.DataPart{ - Data: data, - Metadata: map[string]any{ - a2aDataPartMetaTypeKey: a2aDataPartTypeFunctionCall, - a2aDataPartMetaLongRunningKey: slices.Contains(longRunningToolIDs, part.FunctionCall.ID), - }, - }, nil + r := a2a.NewDataPart(data) + r.SetMeta(a2aDataPartMetaTypeKey, a2aDataPartTypeFunctionCall) + r.SetMeta(a2aDataPartMetaLongRunningKey, slices.Contains(longRunningToolIDs, part.FunctionCall.ID)) + return r, nil } mapStruct, err := converters.ToMapStructure(part) if err != nil { return nil, err } - return a2a.DataPart{Data: mapStruct}, nil + return a2a.NewDataPart(mapStruct), nil } func toGenAIContent(ctx context.Context, msg *a2a.Message, converter A2APartConverter) (*genai.Content, error) { @@ -264,8 +224,8 @@ func toGenAIContent(ctx context.Context, msg *a2a.Message, converter A2APartConv } // ToGenAIPart converts the provided A2A part to a genai equivalent. -func ToGenAIPart(part a2a.Part) (*genai.Part, error) { - parts, err := ToGenAIParts([]a2a.Part{part}) +func ToGenAIPart(part *a2a.Part) (*genai.Part, error) { + parts, err := ToGenAIParts([]*a2a.Part{part}) if err != nil { return nil, err } @@ -273,67 +233,62 @@ func ToGenAIPart(part a2a.Part) (*genai.Part, error) { } // ToGenAIParts converts the provided A2A parts to genai equivalents. -func ToGenAIParts(parts []a2a.Part) ([]*genai.Part, error) { +func ToGenAIParts(parts []*a2a.Part) ([]*genai.Part, error) { result := make([]*genai.Part, len(parts)) for i, part := range parts { - switch v := part.(type) { - case a2a.TextPart: - r := genai.NewPartFromText(v.Text) - if v.Metadata != nil { - if thought, ok := v.Metadata[ToA2AMetaKey("thought")].(bool); ok { - r.Thought = thought - } + if text := part.Text(); text != "" { + r := genai.NewPartFromText(text) + if thought, ok := part.Meta()[ToA2AMetaKey("thought")].(bool); ok { + r.Thought = thought } result[i] = r - - case a2a.DataPart: - r, err := toGenAIDataPart(v) + } else if data := part.Data(); data != nil { + r, err := toGenAIDataPart(part) if err != nil { return nil, err } result[i] = r - - case a2a.FilePart: - r, err := toGenAIFilePart(v) + } else if raw := part.Raw(); raw != nil { + r, err := toGenAIFilePart(part) if err != nil { return nil, err } result[i] = r - - default: - return nil, fmt.Errorf("unknown part type: %T", v) + } else if url := part.URL(); url != "" { + r, err := toGenAIFilePart(part) + if err != nil { + return nil, err + } + result[i] = r + } else { + return nil, fmt.Errorf("unknown part type: %v", part) } } return result, nil } -func toGenAIFilePart(part a2a.FilePart) (*genai.Part, error) { - switch v := part.File.(type) { - case a2a.FileBytes: - bytes, err := base64.StdEncoding.DecodeString(v.Bytes) - if err != nil { - return nil, err - } - data := &genai.Blob{Data: bytes, MIMEType: v.MimeType, DisplayName: v.Name} +func toGenAIFilePart(part *a2a.Part) (*genai.Part, error) { + if raw := part.Raw(); raw != nil { + data := &genai.Blob{Data: raw, MIMEType: part.MediaType, DisplayName: part.Filename} return &genai.Part{InlineData: data}, nil + } - case a2a.FileURI: - data := &genai.FileData{FileURI: v.URI, MIMEType: v.MimeType, DisplayName: v.Name} + if url := part.URL(); url != "" { + data := &genai.FileData{FileURI: string(url), MIMEType: part.MediaType, DisplayName: part.Filename} return &genai.Part{FileData: data}, nil - - default: - return nil, fmt.Errorf("unknown file content type: %T", v) } -} -func toGenAIDataPart(part a2a.DataPart) (*genai.Part, error) { - adkMetaType := part.Metadata[a2aDataPartMetaTypeKey] + return nil, fmt.Errorf("no file content in part") +} - bytes, err := json.Marshal(part.Data) +func toGenAIDataPart(part *a2a.Part) (*genai.Part, error) { + data := part.Data() + bytes, err := json.Marshal(data) if err != nil { return nil, err } + adkMetaType := part.Metadata[a2aDataPartMetaTypeKey] switch adkMetaType { case a2aDataPartTypeCodeExecResult: var val genai.CodeExecutionResult diff --git a/server/adka2a/parts_test.go b/server/adka2a/v2/parts_test.go similarity index 68% rename from server/adka2a/parts_test.go rename to server/adka2a/v2/parts_test.go index 6479f0296..292fbd679 100644 --- a/server/adka2a/parts_test.go +++ b/server/adka2a/v2/parts_test.go @@ -17,7 +17,7 @@ package adka2a import ( "testing" - "github.com/a2aproject/a2a-go/a2a" + "github.com/a2aproject/a2a-go/v2/a2a" "github.com/google/go-cmp/cmp" "google.golang.org/genai" ) @@ -25,51 +25,59 @@ import ( func TestPartsTwoWayConversion(t *testing.T) { testCases := []struct { name string - a2aPart a2a.Part + a2aPart *a2a.Part genaiPart *genai.Part longRunningFunctionIDs []string }{ { name: "text", - a2aPart: a2a.TextPart{Text: "Hello"}, + a2aPart: a2a.NewTextPart("Hello"), genaiPart: &genai.Part{Text: "Hello"}, }, { - name: "thought", - a2aPart: a2a.TextPart{Text: "Hello", Metadata: map[string]any{ToA2AMetaKey("thought"): true}}, + name: "thought", + a2aPart: func() *a2a.Part { + p := a2a.NewTextPart("Hello") + p.SetMeta(ToA2AMetaKey("thought"), true) + return p + }(), genaiPart: &genai.Part{Text: "Hello", Thought: true}, }, { name: "file uri", - a2aPart: a2a.FilePart{ - File: a2a.FileURI{URI: "ftp://cat.com", FileMeta: a2a.FileMeta{MimeType: "image/jpeg", Name: "cat.jpeg"}}, - }, + a2aPart: func() *a2a.Part { + p := a2a.NewFileURLPart("ftp://cat.com", "image/jpeg") + p.Filename = "cat.jpeg" + return p + }(), genaiPart: &genai.Part{ FileData: &genai.FileData{FileURI: "ftp://cat.com", MIMEType: "image/jpeg", DisplayName: "cat.jpeg"}, }, }, { name: "file bytes", - a2aPart: a2a.FilePart{ - File: a2a.FileBytes{Bytes: "/w==", FileMeta: a2a.FileMeta{MimeType: "image/jpeg", Name: "cat.jpeg"}}, - }, + a2aPart: func() *a2a.Part { + p := a2a.NewRawPart([]byte{0xfF}) + p.MediaType = "image/jpeg" + p.Filename = "cat.jpeg" + return p + }(), genaiPart: &genai.Part{ InlineData: &genai.Blob{Data: []byte{0xfF}, MIMEType: "image/jpeg", DisplayName: "cat.jpeg"}, }, }, { name: "function call", - a2aPart: a2a.DataPart{ - Data: map[string]any{ + a2aPart: func() *a2a.Part { + p := a2a.NewDataPart(map[string]any{ "id": "get_weather", "args": map[string]any{"city": "Warsaw"}, "name": "GetWeather", - }, - Metadata: map[string]any{ - a2aDataPartMetaTypeKey: a2aDataPartTypeFunctionCall, - a2aDataPartMetaLongRunningKey: false, - }, - }, + }) + p.SetMeta(a2aDataPartMetaTypeKey, a2aDataPartTypeFunctionCall) + p.SetMeta(a2aDataPartMetaLongRunningKey, false) + return p + }(), genaiPart: &genai.Part{ FunctionCall: &genai.FunctionCall{ ID: "get_weather", @@ -80,17 +88,16 @@ func TestPartsTwoWayConversion(t *testing.T) { }, { name: "long running function call", - a2aPart: a2a.DataPart{ - Data: map[string]any{ + a2aPart: func() *a2a.Part { + p := a2a.NewDataPart(map[string]any{ "id": "get_weather", "args": map[string]any{"city": "Warsaw"}, "name": "GetWeather", - }, - Metadata: map[string]any{ - a2aDataPartMetaTypeKey: a2aDataPartTypeFunctionCall, - a2aDataPartMetaLongRunningKey: true, - }, - }, + }) + p.SetMeta(a2aDataPartMetaTypeKey, a2aDataPartTypeFunctionCall) + p.SetMeta(a2aDataPartMetaLongRunningKey, true) + return p + }(), genaiPart: &genai.Part{ FunctionCall: &genai.FunctionCall{ ID: "get_weather", @@ -102,15 +109,16 @@ func TestPartsTwoWayConversion(t *testing.T) { }, { name: "function response", - a2aPart: a2a.DataPart{ - Data: map[string]any{ + a2aPart: func() *a2a.Part { + p := a2a.NewDataPart(map[string]any{ "id": "get_weather", "scheduling": string(genai.FunctionResponseSchedulingInterrupt), "response": map[string]any{"temperature": "7C"}, "name": "GetWeather", - }, - Metadata: map[string]any{a2aDataPartMetaTypeKey: a2aDataPartTypeFunctionResponse}, - }, + }) + p.SetMeta(a2aDataPartMetaTypeKey, a2aDataPartTypeFunctionResponse) + return p + }(), genaiPart: &genai.Part{ FunctionResponse: &genai.FunctionResponse{ ID: "get_weather", @@ -122,10 +130,11 @@ func TestPartsTwoWayConversion(t *testing.T) { }, { name: "code execution result", - a2aPart: a2a.DataPart{ - Data: map[string]any{"outcome": string(genai.OutcomeOK), "output": "4"}, - Metadata: map[string]any{a2aDataPartMetaTypeKey: a2aDataPartTypeCodeExecResult}, - }, + a2aPart: func() *a2a.Part { + p := a2a.NewDataPart(map[string]any{"outcome": string(genai.OutcomeOK), "output": "4"}) + p.SetMeta(a2aDataPartMetaTypeKey, a2aDataPartTypeCodeExecResult) + return p + }(), genaiPart: &genai.Part{ CodeExecutionResult: &genai.CodeExecutionResult{ Outcome: genai.OutcomeOK, @@ -135,10 +144,11 @@ func TestPartsTwoWayConversion(t *testing.T) { }, { name: "code execution result", - a2aPart: a2a.DataPart{ - Data: map[string]any{"code": "print(2+2)", "language": string(genai.LanguagePython)}, - Metadata: map[string]any{a2aDataPartMetaTypeKey: a2aDataPartTypeCodeExecutableCode}, - }, + a2aPart: func() *a2a.Part { + p := a2a.NewDataPart(map[string]any{"code": "print(2+2)", "language": string(genai.LanguagePython)}) + p.SetMeta(a2aDataPartMetaTypeKey, a2aDataPartTypeCodeExecutableCode) + return p + }(), genaiPart: &genai.Part{ ExecutableCode: &genai.ExecutableCode{ Code: "print(2+2)", @@ -154,11 +164,11 @@ func TestPartsTwoWayConversion(t *testing.T) { if err != nil { t.Errorf("toA2AParts() error = %v, want nil", err) } - if diff := cmp.Diff([]a2a.Part{tc.a2aPart}, toA2A); diff != "" { + if diff := cmp.Diff([]*a2a.Part{tc.a2aPart}, toA2A); diff != "" { t.Errorf("toA2AParts() wrong result (+got,-want)\ngot = %v\nwant = %v\ndiff = %s", toA2A, tc.a2aPart, diff) } - toGenAI, err := ToGenAIParts([]a2a.Part{tc.a2aPart}) + toGenAI, err := ToGenAIParts([]*a2a.Part{tc.a2aPart}) if err != nil { t.Errorf("toGenAIParts() error = %v, want nil", err) } @@ -170,10 +180,10 @@ func TestPartsTwoWayConversion(t *testing.T) { } func TestPartsDataPartConversionRoundTrip(t *testing.T) { - a2aPart := a2a.DataPart{Data: map[string]any{"arbitrary": "data"}} + a2aPart := a2a.NewDataPart(map[string]any{"arbitrary": "data"}) wantGenAI := &genai.Part{InlineData: &genai.Blob{Data: []byte("{\"arbitrary\":\"data\"}"), MIMEType: "text/plain"}} - gotGenAI, err := ToGenAIParts([]a2a.Part{a2aPart}) + gotGenAI, err := ToGenAIParts([]*a2a.Part{a2aPart}) if err != nil { t.Fatalf("toGenAI() error = %v, want nil", err) } @@ -185,7 +195,7 @@ func TestPartsDataPartConversionRoundTrip(t *testing.T) { if err != nil { t.Fatalf("toA2AParts() error = %v, want nil", err) } - if diff := cmp.Diff([]a2a.Part{a2aPart}, gotbackA2A); diff != "" { + if diff := cmp.Diff([]*a2a.Part{a2aPart}, gotbackA2A); diff != "" { t.Fatalf("toA2AParts() wrong result (+got,-want)\ngot = %v\nwant = %v\ndiff = %s", gotbackA2A, a2aPart, diff) } } diff --git a/server/adka2a/processor.go b/server/adka2a/v2/processor.go similarity index 83% rename from server/adka2a/processor.go rename to server/adka2a/v2/processor.go index 696d46bce..94dd4e9ab 100644 --- a/server/adka2a/processor.go +++ b/server/adka2a/v2/processor.go @@ -19,20 +19,21 @@ import ( "fmt" "maps" - "github.com/a2aproject/a2a-go/a2a" - "github.com/a2aproject/a2a-go/a2asrv" + "github.com/a2aproject/a2a-go/v2/a2a" + "github.com/a2aproject/a2a-go/v2/a2asrv" + "github.com/a2aproject/a2a-go/v2/log" "google.golang.org/adk/model" "google.golang.org/adk/session" ) type eventToArtifactTransform interface { - transform(event *session.Event, parts []a2a.Part, meta map[string]any) (*a2a.TaskArtifactUpdateEvent, error) + transform(event *session.Event, parts []*a2a.Part, meta map[string]any) (*a2a.TaskArtifactUpdateEvent, error) makeFinalUpdate() *a2a.TaskArtifactUpdateEvent } type eventProcessor struct { - reqCtx *a2asrv.RequestContext + reqCtx *a2asrv.ExecutorContext meta invocationMeta partConverter GenAIPartConverter @@ -54,7 +55,7 @@ type eventProcessor struct { } func newEventProcessor( - reqCtx *a2asrv.RequestContext, + reqCtx *a2asrv.ExecutorContext, meta invocationMeta, converter GenAIPartConverter, transform eventToArtifactTransform, @@ -80,6 +81,16 @@ func (p *eventProcessor) process(ctx context.Context, event *session.Event) (*a2 return nil, err } + event, err = p.inputRequiredProcessor.process(ctx, event) + if err != nil { + return nil, fmt.Errorf("input required processing failed: %w", err) + } + + parts, err := p.convertParts(ctx, event) + if err != nil { + return nil, err + } + resp := event.LLMResponse if resp.ErrorCode != "" || resp.ErrorMessage != "" { // TODO(yarolegovich): consider merging responses if multiple errors can be produced during an invocation @@ -88,22 +99,13 @@ func (p *eventProcessor) process(ctx context.Context, event *session.Event) (*a2 // not be reflected in this event's metadata terminalEventMeta := maps.Clone(eventMeta) p.failedEvent = toTaskFailedUpdateEvent(p.reqCtx, errorFromResponse(&resp), terminalEventMeta) + p.failedEvent.Status.Message.Parts = append(p.failedEvent.Status.Message.Parts, parts...) } } - event, err = p.inputRequiredProcessor.process(ctx, event) - if err != nil { - return nil, fmt.Errorf("input required processing failed: %w", err) - } - - parts, err := p.convertParts(ctx, event) - if err != nil { - return nil, err - } if len(parts) == 0 { return nil, nil } - result, err := p.eventToArtifact.transform(event, parts, eventMeta) if err != nil { return nil, err @@ -125,7 +127,6 @@ func (p *eventProcessor) makeFinalStatusUpdate() *a2a.TaskStatusUpdateEvent { } ev := a2a.NewStatusUpdateEvent(p.reqCtx, a2a.TaskStateCompleted, nil) - ev.Final = true // we're modifying base processor metadata which might have been sent with one of the previous events. // this update shouldn't be reflected in the sent events' metadata. baseMetaCopy := maps.Clone(p.meta.eventMeta) @@ -133,11 +134,11 @@ func (p *eventProcessor) makeFinalStatusUpdate() *a2a.TaskStatusUpdateEvent { return ev } -func (p *eventProcessor) makeTaskFailedEvent(cause error, event *session.Event) *a2a.TaskStatusUpdateEvent { +func (p *eventProcessor) makeTaskFailedEvent(ctx context.Context, cause error, event *session.Event) *a2a.TaskStatusUpdateEvent { meta := p.meta.eventMeta if event != nil { if eventMeta, err := toEventMeta(p.meta, event); err != nil { - // TODO(yarolegovich): log ignored error + log.Warn(ctx, "failed to convert event metadata for task failed event", "cause", err) } else { meta = eventMeta } @@ -152,7 +153,7 @@ func (p *eventProcessor) updateTerminalActions(event *session.Event) { } } -func (p *eventProcessor) convertParts(ctx context.Context, event *session.Event) ([]a2a.Part, error) { +func (p *eventProcessor) convertParts(ctx context.Context, event *session.Event) ([]*a2a.Part, error) { if event.Content == nil || len(event.Content.Parts) == 0 { return nil, nil } @@ -160,7 +161,7 @@ func (p *eventProcessor) convertParts(ctx context.Context, event *session.Event) if p.partConverter == nil { return ToA2AParts(parts, event.LongRunningToolIDs) } - converted := make([]a2a.Part, 0, len(parts)) + converted := make([]*a2a.Part, 0, len(parts)) for _, part := range parts { cp, err := p.partConverter(ctx, event, part) if err != nil { @@ -175,13 +176,17 @@ func (p *eventProcessor) convertParts(ctx context.Context, event *session.Event) } func toTaskFailedUpdateEvent(task a2a.TaskInfoProvider, cause error, meta map[string]any) *a2a.TaskStatusUpdateEvent { - msg := a2a.NewMessageForTask(a2a.MessageRoleAgent, task, a2a.TextPart{Text: cause.Error()}) + msgPart := a2a.NewTextPart(cause.Error()) + msgPart.Metadata = map[string]any{metadataIsErrMessageKey: true} + msg := a2a.NewMessageForTask(a2a.MessageRoleAgent, task, msgPart) ev := a2a.NewStatusUpdateEvent(task, a2a.TaskStateFailed, msg) ev.Metadata = meta - ev.Final = true return ev } func errorFromResponse(resp *model.LLMResponse) error { - return fmt.Errorf("llm error response: %q", resp.ErrorMessage) + if resp.ErrorCode == "" { + return fmt.Errorf("llm error response: %q", resp.ErrorMessage) + } + return fmt.Errorf("llm error response (code %s): %q", resp.ErrorCode, resp.ErrorMessage) } diff --git a/server/adka2a/processor_test.go b/server/adka2a/v2/processor_test.go similarity index 78% rename from server/adka2a/processor_test.go rename to server/adka2a/v2/processor_test.go index 138bb68bc..d7845996b 100644 --- a/server/adka2a/processor_test.go +++ b/server/adka2a/v2/processor_test.go @@ -17,8 +17,8 @@ package adka2a import ( "testing" - "github.com/a2aproject/a2a-go/a2a" - "github.com/a2aproject/a2a-go/a2asrv" + "github.com/a2aproject/a2a-go/v2/a2a" + "github.com/a2aproject/a2a-go/v2/a2asrv" "github.com/google/go-cmp/cmp" "github.com/google/go-cmp/cmp/cmpopts" "google.golang.org/genai" @@ -37,7 +37,7 @@ func modelPartialResponseFromParts(parts ...*genai.Part) model.LLMResponse { return resp } -func newNonPartialArtifactEvent(task *a2a.Task, parts ...a2a.Part) *a2a.TaskArtifactUpdateEvent { +func newNonPartialArtifactEvent(task *a2a.Task, parts ...*a2a.Part) *a2a.TaskArtifactUpdateEvent { ev := a2a.NewArtifactEvent(task, parts...) // It is important for events to be explicitely marked as ADK partial or non-partial. // This signals to consumers that the remote agent is running its own aggregation logic. @@ -45,22 +45,20 @@ func newNonPartialArtifactEvent(task *a2a.Task, parts ...a2a.Part) *a2a.TaskArti return ev } -func newNonPartialArtifactUpdateEvent(task *a2a.Task, parts ...a2a.Part) *a2a.TaskArtifactUpdateEvent { +func newNonPartialArtifactUpdateEvent(task *a2a.Task, parts ...*a2a.Part) *a2a.TaskArtifactUpdateEvent { ev := newNonPartialArtifactEvent(task, parts...) ev.Append = true return ev } func newDiscardPartialArtifactUpdate(task *a2a.Task) *a2a.TaskArtifactUpdateEvent { - ev := newLegacyPartialArtifactUpdate(task, "", []a2a.Part{a2a.DataPart{Data: map[string]any{}}}) + ev := newLegacyPartialArtifactUpdate(task, "", []*a2a.Part{a2a.NewDataPart(map[string]any{})}) ev.LastChunk = true return ev } func newFinalStatusUpdate(task *a2a.Task, state a2a.TaskState, msg *a2a.Message) *a2a.TaskStatusUpdateEvent { - ev := a2a.NewStatusUpdateEvent(task, state, msg) - ev.Final = true - return ev + return a2a.NewStatusUpdateEvent(task, state, msg) } func TestEventProcessor_Process(t *testing.T) { @@ -95,7 +93,7 @@ func TestEventProcessor_Process(t *testing.T) { LLMResponse: modelResponseFromParts(genai.NewPartFromText("Hello"), genai.NewPartFromText(", world!")), }}, processed: []*a2a.TaskArtifactUpdateEvent{ - newNonPartialArtifactEvent(task, a2a.TextPart{Text: "Hello"}, a2a.TextPart{Text: ", world!"}), + newNonPartialArtifactEvent(task, a2a.NewTextPart("Hello"), a2a.NewTextPart(", world!")), }, terminal: []a2a.Event{newFinalStatusUpdate(task, a2a.TaskStateCompleted, nil)}, }, @@ -107,15 +105,17 @@ func TestEventProcessor_Process(t *testing.T) { {LLMResponse: modelResponseFromParts(genai.NewPartFromText("The answer is 42"))}, }, processed: []*a2a.TaskArtifactUpdateEvent{ - newNonPartialArtifactEvent(task, a2a.DataPart{ - Data: map[string]any{"code": "get_the_answer()", "language": string(genai.LanguagePython)}, - Metadata: map[string]any{a2aDataPartMetaTypeKey: a2aDataPartTypeCodeExecutableCode}, - }), - newNonPartialArtifactUpdateEvent(task, a2a.DataPart{ - Data: map[string]any{"outcome": string(genai.OutcomeOK), "output": "42"}, - Metadata: map[string]any{a2aDataPartMetaTypeKey: a2aDataPartTypeCodeExecResult}, - }), - newNonPartialArtifactUpdateEvent(task, a2a.TextPart{Text: "The answer is 42"}), + newNonPartialArtifactEvent(task, func() *a2a.Part { + p := a2a.NewDataPart(map[string]any{"code": "get_the_answer()", "language": string(genai.LanguagePython)}) + p.SetMeta(a2aDataPartMetaTypeKey, a2aDataPartTypeCodeExecutableCode) + return p + }()), + newNonPartialArtifactUpdateEvent(task, func() *a2a.Part { + p := a2a.NewDataPart(map[string]any{"outcome": string(genai.OutcomeOK), "output": "42"}) + p.SetMeta(a2aDataPartMetaTypeKey, a2aDataPartTypeCodeExecResult) + return p + }()), + newNonPartialArtifactUpdateEvent(task, a2a.NewTextPart("The answer is 42")), }, terminal: []a2a.Event{newFinalStatusUpdate(task, a2a.TaskStateCompleted, nil)}, }, @@ -152,8 +152,8 @@ func TestEventProcessor_Process(t *testing.T) { {LLMResponse: model.LLMResponse{ErrorCode: "1", ErrorMessage: "failed"}}, }, processed: []*a2a.TaskArtifactUpdateEvent{ - newNonPartialArtifactEvent(task, a2a.TextPart{Text: "The answer is"}), - newNonPartialArtifactUpdateEvent(task, a2a.TextPart{Text: "42"}), + newNonPartialArtifactEvent(task, a2a.NewTextPart("The answer is")), + newNonPartialArtifactUpdateEvent(task, a2a.NewTextPart("42")), }, terminal: []a2a.Event{ toTaskFailedUpdateEvent( @@ -170,8 +170,8 @@ func TestEventProcessor_Process(t *testing.T) { {LLMResponse: modelResponseFromParts(genai.NewPartFromText("42"))}, }, processed: []*a2a.TaskArtifactUpdateEvent{ - newNonPartialArtifactEvent(task, a2a.TextPart{Text: "The answer is"}), - newNonPartialArtifactUpdateEvent(task, a2a.TextPart{Text: "42"}), + newNonPartialArtifactEvent(task, a2a.NewTextPart("The answer is")), + newNonPartialArtifactUpdateEvent(task, a2a.NewTextPart("42")), }, terminal: []a2a.Event{ toTaskFailedUpdateEvent( @@ -186,13 +186,12 @@ func TestEventProcessor_Process(t *testing.T) { {LLMResponse: modelResponseFromParts(genai.NewPartFromFunctionCall("get_weather", map[string]any{"city": "Warsaw"}))}, }, processed: []*a2a.TaskArtifactUpdateEvent{ - newNonPartialArtifactEvent(task, a2a.DataPart{ - Data: map[string]any{"name": "get_weather", "args": map[string]any{"city": "Warsaw"}}, - Metadata: map[string]any{ - a2aDataPartMetaTypeKey: a2aDataPartTypeFunctionCall, - a2aDataPartMetaLongRunningKey: false, - }, - }), + newNonPartialArtifactEvent(task, func() *a2a.Part { + p := a2a.NewDataPart(map[string]any{"name": "get_weather", "args": map[string]any{"city": "Warsaw"}}) + p.SetMeta(a2aDataPartMetaTypeKey, a2aDataPartTypeFunctionCall) + p.SetMeta(a2aDataPartMetaLongRunningKey, false) + return p + }()), }, terminal: []a2a.Event{ newFinalStatusUpdate(task, a2a.TaskStateCompleted, nil), @@ -212,14 +211,13 @@ func TestEventProcessor_Process(t *testing.T) { terminal: []a2a.Event{ newFinalStatusUpdate(task, a2a.TaskStateInputRequired, &a2a.Message{ Role: a2a.MessageRoleAgent, - Parts: []a2a.Part{ - a2a.DataPart{ - Data: map[string]any{"id": "get_weather", "name": "weather", "args": map[string]any{"city": "Warsaw"}}, - Metadata: map[string]any{ - a2aDataPartMetaTypeKey: a2aDataPartTypeFunctionCall, - a2aDataPartMetaLongRunningKey: true, - }, - }, + Parts: []*a2a.Part{ + func() *a2a.Part { + p := a2a.NewDataPart(map[string]any{"id": "get_weather", "name": "weather", "args": map[string]any{"city": "Warsaw"}}) + p.SetMeta(a2aDataPartMetaTypeKey, a2aDataPartTypeFunctionCall) + p.SetMeta(a2aDataPartMetaLongRunningKey, true) + return p + }(), }, }), }, @@ -238,20 +236,19 @@ func TestEventProcessor_Process(t *testing.T) { }, }, processed: []*a2a.TaskArtifactUpdateEvent{ - newNonPartialArtifactEvent(task, a2a.TextPart{Text: "This will take a while"}), + newNonPartialArtifactEvent(task, a2a.NewTextPart("This will take a while")), }, terminal: []a2a.Event{ newFinalStatusUpdate(task, a2a.TaskStateInputRequired, &a2a.Message{ Role: a2a.MessageRoleAgent, - Parts: []a2a.Part{ - a2a.DataPart{ - Data: map[string]any{"id": "get_weather", "name": "weather", "args": map[string]any{"city": "Warsaw"}}, - Metadata: map[string]any{ - a2aDataPartMetaTypeKey: a2aDataPartTypeFunctionCall, - a2aDataPartMetaLongRunningKey: true, - }, - }, + Parts: []*a2a.Part{ + func() *a2a.Part { + p := a2a.NewDataPart(map[string]any{"id": "get_weather", "name": "weather", "args": map[string]any{"city": "Warsaw"}}) + p.SetMeta(a2aDataPartMetaTypeKey, a2aDataPartTypeFunctionCall) + p.SetMeta(a2aDataPartMetaLongRunningKey, true) + return p + }(), }, }), }, @@ -276,20 +273,18 @@ func TestEventProcessor_Process(t *testing.T) { terminal: []a2a.Event{ newFinalStatusUpdate(task, a2a.TaskStateInputRequired, &a2a.Message{ Role: a2a.MessageRoleAgent, - Parts: []a2a.Part{ - a2a.DataPart{ - Data: map[string]any{"id": "get_weather", "name": "weather", "args": map[string]any{"city": "Warsaw"}}, - Metadata: map[string]any{ - a2aDataPartMetaTypeKey: a2aDataPartTypeFunctionCall, - a2aDataPartMetaLongRunningKey: true, - }, - }, - a2a.DataPart{ - Data: map[string]any{"id": "get_weather", "name": "weather", "response": map[string]any{"status": "pending"}}, - Metadata: map[string]any{ - a2aDataPartMetaTypeKey: a2aDataPartTypeFunctionResponse, - }, - }, + Parts: []*a2a.Part{ + func() *a2a.Part { + p := a2a.NewDataPart(map[string]any{"id": "get_weather", "name": "weather", "args": map[string]any{"city": "Warsaw"}}) + p.SetMeta(a2aDataPartMetaTypeKey, a2aDataPartTypeFunctionCall) + p.SetMeta(a2aDataPartMetaLongRunningKey, true) + return p + }(), + func() *a2a.Part { + p := a2a.NewDataPart(map[string]any{"id": "get_weather", "name": "weather", "response": map[string]any{"status": "pending"}}) + p.SetMeta(a2aDataPartMetaTypeKey, a2aDataPartTypeFunctionResponse) + return p + }(), }, }), }, @@ -305,7 +300,6 @@ func TestEventProcessor_Process(t *testing.T) { ContextID: task.ContextID, Status: a2a.TaskStatus{State: a2a.TaskStateCompleted}, Metadata: map[string]any{metadataEscalateKey: true, metadataTransferToAgentKey: "a-2"}, - Final: true, }, }, }, @@ -321,7 +315,6 @@ func TestEventProcessor_Process(t *testing.T) { ContextID: task.ContextID, Status: a2a.TaskStatus{State: a2a.TaskStateCompleted}, Metadata: map[string]any{metadataTransferToAgentKey: "a-3"}, - Final: true, }, }, }, @@ -335,7 +328,7 @@ func TestEventProcessor_Process(t *testing.T) { {LLMResponse: model.LLMResponse{ErrorCode: "1", ErrorMessage: "failed"}}, }, processed: []*a2a.TaskArtifactUpdateEvent{ - newNonPartialArtifactEvent(task, a2a.TextPart{Text: "The answer is"}), + newNonPartialArtifactEvent(task, a2a.NewTextPart("The answer is")), }, terminal: []a2a.Event{ toTaskFailedUpdateEvent( @@ -358,25 +351,31 @@ func TestEventProcessor_Process(t *testing.T) { )}, }, processed: []*a2a.TaskArtifactUpdateEvent{ - newLegacyPartialArtifactUpdate(task, artifactIDPlaceholder, []a2a.Part{ - a2a.TextPart{Text: "The answer is", Metadata: map[string]any{ToA2AMetaKey("partial"): true}}, - a2a.DataPart{ - Data: map[string]any{"code": "get_the_answer()", "language": string(genai.LanguagePython)}, - Metadata: map[string]any{ - a2aDataPartMetaTypeKey: a2aDataPartTypeCodeExecutableCode, - ToA2AMetaKey("partial"): true, - }, - }, + newLegacyPartialArtifactUpdate(task, artifactIDPlaceholder, []*a2a.Part{ + func() *a2a.Part { + p := a2a.NewTextPart("The answer is") + p.SetMeta(ToA2AMetaKey("partial"), true) + return p + }(), + func() *a2a.Part { + p := a2a.NewDataPart(map[string]any{"code": "get_the_answer()", "language": string(genai.LanguagePython)}) + p.SetMeta(a2aDataPartMetaTypeKey, a2aDataPartTypeCodeExecutableCode) + p.SetMeta(ToA2AMetaKey("partial"), true) + return p + }(), }), - newLegacyPartialArtifactUpdate(task, artifactIDPlaceholder, []a2a.Part{ - a2a.DataPart{ - Data: map[string]any{"outcome": string(genai.OutcomeOK), "output": "42"}, - Metadata: map[string]any{ - a2aDataPartMetaTypeKey: a2aDataPartTypeCodeExecResult, - ToA2AMetaKey("partial"): true, - }, - }, - a2a.TextPart{Text: "42", Metadata: map[string]any{ToA2AMetaKey("partial"): true}}, + newLegacyPartialArtifactUpdate(task, artifactIDPlaceholder, []*a2a.Part{ + func() *a2a.Part { + p := a2a.NewDataPart(map[string]any{"outcome": string(genai.OutcomeOK), "output": "42"}) + p.SetMeta(a2aDataPartMetaTypeKey, a2aDataPartTypeCodeExecResult) + p.SetMeta(ToA2AMetaKey("partial"), true) + return p + }(), + func() *a2a.Part { + p := a2a.NewTextPart("42") + p.SetMeta(ToA2AMetaKey("partial"), true) + return p + }(), }), }, terminal: []a2a.Event{ @@ -392,7 +391,7 @@ func TestEventProcessor_Process(t *testing.T) { }}, processed: []*a2a.TaskArtifactUpdateEvent{ func() *a2a.TaskArtifactUpdateEvent { - ev := newNonPartialArtifactEvent(task, a2a.TextPart{Text: "Hello"}) + ev := newNonPartialArtifactEvent(task, a2a.NewTextPart("Hello")) ev.Metadata[ToA2AMetaKey("invocation_id")] = "test-invocation-id" return ev }(), @@ -408,7 +407,7 @@ func TestEventProcessor_Process(t *testing.T) { cmpopts.IgnoreFields(a2a.TaskStatus{}, "Timestamp"), } t.Run(tc.name, func(t *testing.T) { - reqCtx := &a2asrv.RequestContext{TaskID: task.ID, ContextID: task.ContextID} + reqCtx := &a2asrv.ExecutorContext{TaskID: task.ID, ContextID: task.ContextID} processor := newEventProcessor(reqCtx, invocationMeta{}, nil, newLegacyArtifactMaker(reqCtx)) var gotEvents []*a2a.TaskArtifactUpdateEvent @@ -456,7 +455,7 @@ func TestEventProcessor_ArtifactUpdates(t *testing.T) { }, } - reqCtx := &a2asrv.RequestContext{TaskID: task.ID, ContextID: task.ContextID} + reqCtx := &a2asrv.ExecutorContext{TaskID: task.ID, ContextID: task.ContextID} processor := newEventProcessor(reqCtx, invocationMeta{}, nil, newLegacyArtifactMaker(reqCtx)) got := make([]*a2a.TaskArtifactUpdateEvent, len(events)) for i, event := range events { @@ -499,7 +498,7 @@ func TestEventProcessor_PartialEventsAreDiscardedAsAnArtifact(t *testing.T) { {LLMResponse: modelResponseFromParts(genai.NewPartFromText("Hello, world!"))}, } - reqCtx := &a2asrv.RequestContext{TaskID: task.ID, ContextID: task.ContextID} + reqCtx := &a2asrv.ExecutorContext{TaskID: task.ID, ContextID: task.ContextID} processor := newEventProcessor(reqCtx, invocationMeta{}, nil, newLegacyArtifactMaker(reqCtx)) got := make([]*a2a.TaskArtifactUpdateEvent, len(events)) for i, event := range events { @@ -544,8 +543,12 @@ func TestEventProcessor_PartialEventsAreDiscardedAsAnArtifact(t *testing.T) { TaskID: task.ID, ContextID: task.ContextID, Artifact: &a2a.Artifact{ - ID: got[0].Artifact.ID, - Parts: a2a.ContentParts{a2a.DataPart{Data: map[string]any{}, Metadata: map[string]any{metadataPartialKey: true}}}, + ID: got[0].Artifact.ID, + Parts: a2a.ContentParts{func() *a2a.Part { + p := a2a.NewDataPart(map[string]any{}) + p.SetMeta(metadataPartialKey, true) + return p + }()}, Metadata: map[string]any{metadataPartialKey: true}, }, Metadata: map[string]any{metadataPartialKey: true}, diff --git a/server/adka2a/task_artifact.go b/server/adka2a/v2/task_artifact.go similarity index 86% rename from server/adka2a/task_artifact.go rename to server/adka2a/v2/task_artifact.go index 4458f8a20..6200b79d4 100644 --- a/server/adka2a/task_artifact.go +++ b/server/adka2a/v2/task_artifact.go @@ -17,18 +17,18 @@ package adka2a import ( "maps" - "github.com/a2aproject/a2a-go/a2a" - "github.com/a2aproject/a2a-go/a2asrv" + "github.com/a2aproject/a2a-go/v2/a2a" + "github.com/a2aproject/a2a-go/v2/a2asrv" "google.golang.org/adk/session" ) type artifactMaker struct { - reqCtx *a2asrv.RequestContext + reqCtx *a2asrv.ExecutorContext lastAgentPartialArtifact map[string]a2a.ArtifactID } -func newArtifactMaker(reqCtx *a2asrv.RequestContext) *artifactMaker { +func newArtifactMaker(reqCtx *a2asrv.ExecutorContext) *artifactMaker { return &artifactMaker{ reqCtx: reqCtx, lastAgentPartialArtifact: make(map[string]a2a.ArtifactID), @@ -37,7 +37,7 @@ func newArtifactMaker(reqCtx *a2asrv.RequestContext) *artifactMaker { var _ eventToArtifactTransform = (*artifactMaker)(nil) -func (m *artifactMaker) transform(event *session.Event, parts []a2a.Part, meta map[string]any) (*a2a.TaskArtifactUpdateEvent, error) { +func (m *artifactMaker) transform(event *session.Event, parts []*a2a.Part, meta map[string]any) (*a2a.TaskArtifactUpdateEvent, error) { result := a2a.NewArtifactEvent(m.reqCtx, parts...) if artifactID, ok := m.lastAgentPartialArtifact[event.Author]; ok { @@ -68,7 +68,7 @@ func (m *artifactMaker) makeFinalUpdate() *a2a.TaskArtifactUpdateEvent { } type legacyArtifactMaker struct { - reqCtx *a2asrv.RequestContext + reqCtx *a2asrv.ExecutorContext // responseID is created once the first TaskArtifactUpdateEvent is sent. Used for subsequent artifact updates. responseID a2a.ArtifactID @@ -79,7 +79,7 @@ type legacyArtifactMaker struct { partialResponseID a2a.ArtifactID } -func newLegacyArtifactMaker(reqCtx *a2asrv.RequestContext) *legacyArtifactMaker { +func newLegacyArtifactMaker(reqCtx *a2asrv.ExecutorContext) *legacyArtifactMaker { return &legacyArtifactMaker{ reqCtx: reqCtx, } @@ -87,7 +87,7 @@ func newLegacyArtifactMaker(reqCtx *a2asrv.RequestContext) *legacyArtifactMaker var _ eventToArtifactTransform = (*legacyArtifactMaker)(nil) -func (p *legacyArtifactMaker) transform(event *session.Event, parts []a2a.Part, meta map[string]any) (*a2a.TaskArtifactUpdateEvent, error) { +func (p *legacyArtifactMaker) transform(event *session.Event, parts []*a2a.Part, meta map[string]any) (*a2a.TaskArtifactUpdateEvent, error) { var result *a2a.TaskArtifactUpdateEvent if event.Partial { result = newLegacyPartialArtifactUpdate(p.reqCtx, p.partialResponseID, parts) @@ -109,12 +109,12 @@ func (p *legacyArtifactMaker) makeFinalUpdate() *a2a.TaskArtifactUpdateEvent { if p.partialResponseID == "" { return nil } - ev := newLegacyPartialArtifactUpdate(p.reqCtx, p.partialResponseID, []a2a.Part{a2a.DataPart{Data: map[string]any{}}}) + ev := newLegacyPartialArtifactUpdate(p.reqCtx, p.partialResponseID, []*a2a.Part{a2a.NewDataPart(map[string]any{})}) ev.LastChunk = true return ev } -func newLegacyArtifactUpdate(task a2a.TaskInfoProvider, id a2a.ArtifactID, parts []a2a.Part) *a2a.TaskArtifactUpdateEvent { +func newLegacyArtifactUpdate(task a2a.TaskInfoProvider, id a2a.ArtifactID, parts []*a2a.Part) *a2a.TaskArtifactUpdateEvent { var result *a2a.TaskArtifactUpdateEvent if id == "" { result = a2a.NewArtifactEvent(task, parts...) @@ -127,7 +127,7 @@ func newLegacyArtifactUpdate(task a2a.TaskInfoProvider, id a2a.ArtifactID, parts return result } -func newLegacyPartialArtifactUpdate(task a2a.TaskInfoProvider, artifactID a2a.ArtifactID, parts []a2a.Part) *a2a.TaskArtifactUpdateEvent { +func newLegacyPartialArtifactUpdate(task a2a.TaskInfoProvider, artifactID a2a.ArtifactID, parts []*a2a.Part) *a2a.TaskArtifactUpdateEvent { ev := newLegacyArtifactUpdate(task, artifactID, parts) updatePartsMetadata(parts, map[string]any{metadataPartialKey: true}) if ev.Artifact.Metadata == nil { diff --git a/server/adka2a/utils.go b/server/adka2a/v2/utils.go similarity index 97% rename from server/adka2a/utils.go rename to server/adka2a/v2/utils.go index 04f60615d..c8bf31835 100644 --- a/server/adka2a/utils.go +++ b/server/adka2a/v2/utils.go @@ -15,7 +15,7 @@ package adka2a import ( - "github.com/a2aproject/a2a-go/a2a" + "github.com/a2aproject/a2a-go/v2/a2a" "google.golang.org/adk/agent" iagent "google.golang.org/adk/internal/agent" diff --git a/server/adkrest/controllers/runtime.go b/server/adkrest/controllers/runtime.go index 11085eeba..7eaea1119 100644 --- a/server/adkrest/controllers/runtime.go +++ b/server/adkrest/controllers/runtime.go @@ -22,6 +22,9 @@ import ( "net/http" "time" + "github.com/gorilla/websocket" + "google.golang.org/genai" + "google.golang.org/adk/agent" "google.golang.org/adk/artifact" "google.golang.org/adk/memory" @@ -32,17 +35,18 @@ import ( // RuntimeAPIController is the controller for the Runtime API. type RuntimeAPIController struct { - sseTimeout time.Duration - sessionService session.Service - memoryService memory.Service - artifactService artifact.Service - agentLoader agent.Loader - pluginConfig runner.PluginConfig + sseTimeout time.Duration + sessionService session.Service + memoryService memory.Service + artifactService artifact.Service + agentLoader agent.Loader + pluginConfig runner.PluginConfig + autoCreateSession bool } // NewRuntimeAPIController creates the controller for the Runtime API. -func NewRuntimeAPIController(sessionService session.Service, memoryService memory.Service, agentLoader agent.Loader, artifactService artifact.Service, sseTimeout time.Duration, pluginConfig runner.PluginConfig) *RuntimeAPIController { - return &RuntimeAPIController{sessionService: sessionService, memoryService: memoryService, agentLoader: agentLoader, artifactService: artifactService, sseTimeout: sseTimeout, pluginConfig: pluginConfig} +func NewRuntimeAPIController(sessionService session.Service, memoryService memory.Service, agentLoader agent.Loader, artifactService artifact.Service, sseTimeout time.Duration, pluginConfig runner.PluginConfig, autoCreateSession bool) *RuntimeAPIController { + return &RuntimeAPIController{sessionService: sessionService, memoryService: memoryService, agentLoader: agentLoader, artifactService: artifactService, sseTimeout: sseTimeout, pluginConfig: pluginConfig, autoCreateSession: autoCreateSession} } // RunAgent executes a non-streaming agent run for a given session and message. @@ -75,7 +79,11 @@ func (c *RuntimeAPIController) runAgent(ctx context.Context, runAgentRequest mod return nil, err } - resp := r.Run(ctx, runAgentRequest.UserId, runAgentRequest.SessionId, &runAgentRequest.NewMessage, *rCfg) + var opts []runner.RunOption + if runAgentRequest.StateDelta != nil { + opts = append(opts, runner.WithStateDelta(*runAgentRequest.StateDelta)) + } + resp := r.Run(ctx, runAgentRequest.UserId, runAgentRequest.SessionId, &runAgentRequest.NewMessage, *rCfg, opts...) var events []*session.Event for event, err := range resp { @@ -204,12 +212,13 @@ func (c *RuntimeAPIController) getRunner(req models.RunAgentRequest) (*runner.Ru } r, err := runner.New(runner.Config{ - AppName: req.AppName, - Agent: curAgent, - SessionService: c.sessionService, - MemoryService: c.memoryService, - ArtifactService: c.artifactService, - PluginConfig: c.pluginConfig, + AppName: req.AppName, + Agent: curAgent, + SessionService: c.sessionService, + MemoryService: c.memoryService, + ArtifactService: c.artifactService, + PluginConfig: c.pluginConfig, + AutoCreateSession: c.autoCreateSession, }, ) if err != nil { @@ -225,11 +234,8 @@ func (c *RuntimeAPIController) getRunner(req models.RunAgentRequest) (*runner.Ru }, nil } -func decodeRequestBody(req *http.Request) (decodedReq models.RunAgentRequest, err error) { +func decodeRequestBody(req *http.Request) (models.RunAgentRequest, error) { var runAgentRequest models.RunAgentRequest - defer func() { - err = req.Body.Close() - }() d := json.NewDecoder(req.Body) d.DisallowUnknownFields() if err := d.Decode(&runAgentRequest); err != nil { @@ -237,3 +243,146 @@ func decodeRequestBody(req *http.Request) (decodedReq models.RunAgentRequest, er } return runAgentRequest, nil } + +func (c *RuntimeAPIController) RunLiveHandler(rw http.ResponseWriter, req *http.Request) error { + upgrader := websocket.Upgrader{ + ReadBufferSize: 1024, + WriteBufferSize: 1024, + } + + q := req.URL.Query() + appName := q.Get("appName") + if appName == "" { + appName = q.Get("app_name") + } + userID := q.Get("userId") + if userID == "" { + userID = q.Get("user_id") + } + sessionID := q.Get("sessionId") + if sessionID == "" { + sessionID = q.Get("session_id") + } + + if appName == "" || userID == "" || sessionID == "" { + return fmt.Errorf("appName, userId, and sessionId are required") + } + + ws, err := upgrader.Upgrade(rw, req, nil) + if err != nil { + return fmt.Errorf("failed to upgrade to websocket: %w", err) + } + defer func() { + _ = ws.Close() + }() + + sendClose := func(code int, reason string) { + _ = ws.WriteMessage(websocket.CloseMessage, websocket.FormatCloseMessage(code, reason)) + _ = ws.SetReadDeadline(time.Now().Add(time.Second)) + for { + if _, _, err := ws.ReadMessage(); err != nil { + break + } + } + } + + r, _, err := c.getRunner(models.RunAgentRequest{AppName: appName, UserId: userID, SessionId: sessionID}) + if err != nil { + closeReason := err.Error() + if _, loadErr := c.agentLoader.LoadAgent(appName); loadErr != nil { + closeReason = fmt.Sprintf("agent %s not found for original error: %v", appName, err) + } + log.Printf("Failed to get runner for app %s: %v", appName, err) + sendClose(websocket.CloseInternalServerErr, closeReason) + return nil + } + + // Read from Runner and write back to client over the WebSocket + liveSession, eventIter, err := r.RunLive(req.Context(), userID, sessionID, agent.LiveRunConfig{ + MaxLLMCalls: 100, // Reasonable default + ResponseModalities: []genai.Modality{genai.ModalityAudio}, + InputAudioTranscription: &genai.AudioTranscriptionConfig{}, + OutputAudioTranscription: &genai.AudioTranscriptionConfig{}, + }) + if err != nil { + log.Printf("RunLive failed for app %s: %v", appName, err) + sendClose(websocket.CloseInternalServerErr, err.Error()) + return nil + } + defer func() { + _ = liveSession.Close() + }() + + // Spawning goroutine for reading from the client over WebSocket and pushing it to Runner + go func() { + defer func() { + _ = liveSession.Close() + }() + for { + messageType, p, err := ws.ReadMessage() + if err != nil { + if !websocket.IsCloseError(err, websocket.CloseNormalClosure, websocket.CloseGoingAway) { + log.Printf("WebSocket read error for app %s: %v", appName, err) + } + break + } + + if messageType == websocket.BinaryMessage { + if err := liveSession.Send(agent.LiveRequest{ + RealtimeInput: &genai.Blob{ + MIMEType: "audio/pcm;rate=16000", + Data: p, + }, + }); err != nil { + log.Printf("Failed to send binary data to Gemini for app %s: %v", appName, err) + break + } + } else if messageType == websocket.TextMessage { + var apiReq models.LiveRequest + if err := json.Unmarshal(p, &apiReq); err != nil { + log.Printf("Failed to unmarshal client message for app %s: %v", appName, err) + continue + } + + if apiReq.Close { + break + } + + liveReq := agent.LiveRequest{ + Content: apiReq.Content, + } + + if apiReq.ActivityStart != nil { + liveReq.RealtimeInput = apiReq.ActivityStart + } else if apiReq.ActivityEnd != nil { + liveReq.RealtimeInput = apiReq.ActivityEnd + } else if apiReq.Blob != nil { + liveReq.RealtimeInput = &genai.Blob{ + MIMEType: apiReq.Blob.MIMEType, + Data: apiReq.Blob.Data, + } + } + + if err := liveSession.Send(liveReq); err != nil { + log.Printf("Failed to send message to Gemini for app %s: %v", appName, err) + break + } + } + } + }() + + for event, err := range eventIter { + if err != nil { + log.Printf("RunLive failed: %v\n", err) + _ = ws.WriteMessage(websocket.CloseMessage, websocket.FormatCloseMessage(websocket.CloseInternalServerErr, err.Error())) + break + } + + err = ws.WriteJSON(models.FromSessionEvent(*event)) + if err != nil { + break + } + } + + return nil +} diff --git a/server/adkrest/controllers/runtime_test.go b/server/adkrest/controllers/runtime_test.go index a6672fb83..7172d340b 100644 --- a/server/adkrest/controllers/runtime_test.go +++ b/server/adkrest/controllers/runtime_test.go @@ -16,6 +16,7 @@ package controllers import ( "bytes" + "context" "encoding/json" "fmt" "iter" @@ -77,7 +78,7 @@ func TestNewRuntimeAPIController_PluginsAssignment(t *testing.T) { t.Run(tt.name, func(t *testing.T) { controller := NewRuntimeAPIController(nil, nil, nil, nil, 10*time.Second, runner.PluginConfig{ Plugins: tt.plugins, - }) + }, false) if controller == nil { t.Fatal("NewRuntimeAPIController returned nil") @@ -116,7 +117,7 @@ func testAgent(results []testAgentResult) func(ctx agent.InvocationContext) iter } func makeEvent(id, author, text string) *session.Event { - e := session.NewEvent(id) + e := session.NewEventWithContext(context.Background(), id) e.Author = author e.LLMResponse.Content = &genai.Content{ Parts: []*genai.Part{{Text: text}}, @@ -201,6 +202,7 @@ func TestRunSSEHandler(t *testing.T) { nil, 10*time.Second, runner.PluginConfig{}, + false, ) // Create request diff --git a/server/adkrest/handler.go b/server/adkrest/handler.go index 1bc32d331..e96e9a67e 100644 --- a/server/adkrest/handler.go +++ b/server/adkrest/handler.go @@ -47,7 +47,7 @@ func NewServer(cfg ServerConfig) (*Server, error) { // where the ADK REST API will be served. setupRouter(router, routers.NewSessionsAPIRouter(controllers.NewSessionsAPIController(cfg.SessionService)), - routers.NewRuntimeAPIRouter(controllers.NewRuntimeAPIController(cfg.SessionService, cfg.MemoryService, cfg.AgentLoader, cfg.ArtifactService, cfg.SSEWriteTimeout, cfg.PluginConfig)), + routers.NewRuntimeAPIRouter(controllers.NewRuntimeAPIController(cfg.SessionService, cfg.MemoryService, cfg.AgentLoader, cfg.ArtifactService, cfg.SSEWriteTimeout, cfg.PluginConfig, false)), routers.NewAppsAPIRouter(controllers.NewAppsAPIController(cfg.AgentLoader)), routers.NewDebugAPIRouter(controllers.NewDebugAPIController(cfg.SessionService, cfg.AgentLoader, debugTelemetry)), routers.NewArtifactsAPIRouter(controllers.NewArtifactsAPIController(cfg.ArtifactService)), diff --git a/server/adkrest/internal/models/event.go b/server/adkrest/internal/models/event.go index 79497ff03..6fb42cd61 100644 --- a/server/adkrest/internal/models/event.go +++ b/server/adkrest/internal/models/event.go @@ -34,23 +34,25 @@ type EventActions struct { // Event represents a single event in a session. type Event struct { - ID string `json:"id"` - InvocationID string `json:"invocationId"` - Branch string `json:"branch,omitempty"` - Author string `json:"author"` - Partial bool `json:"partial,omitempty"` - LongRunningToolIDs []string `json:"longRunningToolIds,omitempty"` - Content *genai.Content `json:"content"` - GroundingMetadata *genai.GroundingMetadata `json:"groundingMetadata"` - UsageMetadata *genai.GenerateContentResponseUsageMetadata `json:"usageMetadata"` - TurnComplete bool `json:"turnComplete,omitempty"` - Interrupted bool `json:"interrupted,omitempty"` - ErrorCode string `json:"errorCode,omitempty"` - ErrorMessage string `json:"errorMessage,omitempty"` - AvgLogprobs float64 `json:"avgLogprobs,omitempty"` - FinishReason genai.FinishReason `json:"finishReason,omitempty"` - ModelVersion string `json:"modelVersion,omitempty"` - Actions EventActions `json:"actions"` + ID string `json:"id"` + InvocationID string `json:"invocationId"` + Branch string `json:"branch,omitempty"` + Author string `json:"author"` + Partial bool `json:"partial,omitempty"` + LongRunningToolIDs []string `json:"longRunningToolIds,omitempty"` + Content *genai.Content `json:"content"` + GroundingMetadata *genai.GroundingMetadata `json:"groundingMetadata"` + UsageMetadata *genai.GenerateContentResponseUsageMetadata `json:"usageMetadata"` + TurnComplete bool `json:"turnComplete,omitempty"` + Interrupted bool `json:"interrupted,omitempty"` + ErrorCode string `json:"errorCode,omitempty"` + ErrorMessage string `json:"errorMessage,omitempty"` + AvgLogprobs float64 `json:"avgLogprobs,omitempty"` + FinishReason genai.FinishReason `json:"finishReason,omitempty"` + ModelVersion string `json:"modelVersion,omitempty"` + InputTranscription *genai.Transcription `json:"inputTranscription,omitempty"` + OutputTranscription *genai.Transcription `json:"outputTranscription,omitempty"` + Actions EventActions `json:"actions"` } // ToSessionEvent maps Event data struct to session.Event @@ -62,17 +64,19 @@ func ToSessionEvent(event Event) *session.Event { Author: event.Author, LongRunningToolIDs: event.LongRunningToolIDs, LLMResponse: model.LLMResponse{ - AvgLogprobs: event.AvgLogprobs, - Content: event.Content, - GroundingMetadata: event.GroundingMetadata, - UsageMetadata: event.UsageMetadata, - Partial: event.Partial, - TurnComplete: event.TurnComplete, - Interrupted: event.Interrupted, - ErrorCode: event.ErrorCode, - ErrorMessage: event.ErrorMessage, - FinishReason: event.FinishReason, - ModelVersion: event.ModelVersion, + AvgLogprobs: event.AvgLogprobs, + Content: event.Content, + GroundingMetadata: event.GroundingMetadata, + UsageMetadata: event.UsageMetadata, + Partial: event.Partial, + TurnComplete: event.TurnComplete, + Interrupted: event.Interrupted, + ErrorCode: event.ErrorCode, + ErrorMessage: event.ErrorMessage, + FinishReason: event.FinishReason, + ModelVersion: event.ModelVersion, + InputTranscription: event.InputTranscription, + OutputTranscription: event.OutputTranscription, }, Actions: session.EventActions{ StateDelta: event.Actions.StateDelta, @@ -87,22 +91,24 @@ func ToSessionEvent(event Event) *session.Event { // FromSessionEvent maps session.Event to Event data struct func FromSessionEvent(event session.Event) Event { return Event{ - ID: event.ID, - InvocationID: event.InvocationID, - Branch: event.Branch, - Author: event.Author, - Partial: event.Partial, - LongRunningToolIDs: event.LongRunningToolIDs, - AvgLogprobs: event.LLMResponse.AvgLogprobs, - Content: event.LLMResponse.Content, - GroundingMetadata: event.LLMResponse.GroundingMetadata, - UsageMetadata: event.LLMResponse.UsageMetadata, - TurnComplete: event.LLMResponse.TurnComplete, - Interrupted: event.LLMResponse.Interrupted, - ErrorCode: event.LLMResponse.ErrorCode, - ErrorMessage: event.LLMResponse.ErrorMessage, - FinishReason: event.LLMResponse.FinishReason, - ModelVersion: event.LLMResponse.ModelVersion, + ID: event.ID, + InvocationID: event.InvocationID, + Branch: event.Branch, + Author: event.Author, + Partial: event.Partial, + LongRunningToolIDs: event.LongRunningToolIDs, + AvgLogprobs: event.LLMResponse.AvgLogprobs, + Content: event.LLMResponse.Content, + GroundingMetadata: event.LLMResponse.GroundingMetadata, + UsageMetadata: event.LLMResponse.UsageMetadata, + TurnComplete: event.LLMResponse.TurnComplete, + Interrupted: event.LLMResponse.Interrupted, + ErrorCode: event.LLMResponse.ErrorCode, + ErrorMessage: event.LLMResponse.ErrorMessage, + FinishReason: event.LLMResponse.FinishReason, + ModelVersion: event.LLMResponse.ModelVersion, + InputTranscription: event.LLMResponse.InputTranscription, + OutputTranscription: event.LLMResponse.OutputTranscription, Actions: EventActions{ StateDelta: event.Actions.StateDelta, ArtifactDelta: event.Actions.ArtifactDelta, diff --git a/server/adkrest/internal/models/runtime.go b/server/adkrest/internal/models/runtime.go index 59c82a552..e4888c84c 100644 --- a/server/adkrest/internal/models/runtime.go +++ b/server/adkrest/internal/models/runtime.go @@ -50,3 +50,18 @@ func (req RunAgentRequest) AssertRunAgentRequestRequired() error { return nil } + +// blob represents a genai.blob sent by the client, explicitly mapping mime_type. +type blob struct { + MIMEType string `json:"mime_type,omitempty"` + Data []byte `json:"data,omitempty"` +} + +// LiveRequest represents the client request format for real-time interactions over WebSocket. +type LiveRequest struct { + Content *genai.Content `json:"content,omitempty"` + Blob *blob `json:"blob,omitempty"` + ActivityStart *genai.ActivityStart `json:"activityStart,omitempty"` + ActivityEnd *genai.ActivityEnd `json:"activityEnd,omitempty"` + Close bool `json:"close,omitempty"` +} diff --git a/server/adkrest/internal/routers/runtime.go b/server/adkrest/internal/routers/runtime.go index 3819ce64a..1bb3f0a5d 100644 --- a/server/adkrest/internal/routers/runtime.go +++ b/server/adkrest/internal/routers/runtime.go @@ -45,5 +45,11 @@ func (r *RuntimeAPIRouter) Routes() Routes { Pattern: "/run_sse", HandlerFunc: r.runtimeController.RunSSEHandler, }, + Route{ + Name: "RunAgentLive", + Methods: []string{http.MethodGet, http.MethodOptions}, + Pattern: "/run_live", + HandlerFunc: controllers.NewErrorHandler(r.runtimeController.RunLiveHandler), + }, } } diff --git a/server/adkrest/internal/services/debugtelemetry_test.go b/server/adkrest/internal/services/debugtelemetry_test.go index 3ebb6ed8a..0daeb6ef9 100644 --- a/server/adkrest/internal/services/debugtelemetry_test.go +++ b/server/adkrest/internal/services/debugtelemetry_test.go @@ -228,6 +228,7 @@ func TestDebugTelemetryGetSpansBySessionID(t *testing.T) { cmpopts.IgnoreUnexported(log.Value{}), cmpopts.IgnoreFields(DebugSpan{}, "StartTime", "EndTime", "TraceID", "SpanID", "ParentSpanID"), cmpopts.IgnoreFields(DebugLog{}, "ObservedTimestamp", "TraceID", "SpanID"), + cmpopts.SortSlices(compareDebugSpans), cmpopts.EquateEmpty(), } @@ -365,6 +366,7 @@ func TestDebugTelemetryGetSpansByEventID(t *testing.T) { cmpopts.IgnoreUnexported(log.Value{}), cmpopts.IgnoreFields(DebugSpan{}, "StartTime", "EndTime", "ParentSpanID", "TraceID", "SpanID"), cmpopts.IgnoreFields(DebugLog{}, "ObservedTimestamp", "TraceID", "SpanID"), + cmpopts.SortSlices(compareDebugSpans), cmpopts.EquateEmpty(), } @@ -472,3 +474,19 @@ func setupWithConfig(t *testing.T, cfg *DebugTelemetryConfig) (*DebugTelemetry, func setup(t *testing.T) (*DebugTelemetry, *sdktrace.TracerProvider, *sdklog.LoggerProvider) { return setupWithConfig(t, nil) } + +func compareDebugSpans(a, b DebugSpan) bool { + if a.Name != b.Name { + return a.Name < b.Name + } + if a.ParentSpanID != b.ParentSpanID { + return a.ParentSpanID < b.ParentSpanID + } + if a.Attributes[string(semconv.GenAIConversationIDKey)] != b.Attributes[string(semconv.GenAIConversationIDKey)] { + return a.Attributes[string(semconv.GenAIConversationIDKey)] < b.Attributes[string(semconv.GenAIConversationIDKey)] + } + if a.Attributes[eventIDKey] != b.Attributes[eventIDKey] { + return a.Attributes[eventIDKey] < b.Attributes[eventIDKey] + } + return a.Attributes["genai.operation.name"] < b.Attributes["genai.operation.name"] +} diff --git a/server/agentengine/controllers/method/create_session.go b/server/agentengine/controllers/method/create_session.go index bc27f37e3..b49d942ec 100644 --- a/server/agentengine/controllers/method/create_session.go +++ b/server/agentengine/controllers/method/create_session.go @@ -45,10 +45,6 @@ func (c *createSessionHandler) Metadata() (*structpb.Struct, error) { "name": c.methodName, "parameters": map[string]any{ "properties": map[string]any{ - "session_id": map[string]any{ - "nullable": true, - "type": "string", - }, "user_id": map[string]any{ "type": "string", }, @@ -67,8 +63,6 @@ func (c *createSessionHandler) Metadata() (*structpb.Struct, error) { Args: user_id (str): Required. The ID of the user. - session_id (str): - Optional. The ID of the session. If not provided, an ID will be generated for the session. state (dict[str, Any]): Optional. The initial state of the session. @@ -92,10 +86,9 @@ func (c *createSessionHandler) Handle(ctx context.Context, rw http.ResponseWrite } ssReq := &session.CreateRequest{ - AppName: c.agentEngineID, - UserID: req.Input.UserID, - SessionID: req.Input.SessionID, - State: req.Input.State, + AppName: c.agentEngineID, + UserID: req.Input.UserID, + State: req.Input.State, } resp, err := c.sessionService.Create(ctx, ssReq) if err != nil { diff --git a/server/agentengine/controllers/method/delete_session.go b/server/agentengine/controllers/method/delete_session.go index 246de6bc7..3c4d20e20 100644 --- a/server/agentengine/controllers/method/delete_session.go +++ b/server/agentengine/controllers/method/delete_session.go @@ -92,14 +92,11 @@ func (g *deleteSessionHandler) Handle(ctx context.Context, rw http.ResponseWrite SessionID: req.Input.SessionID, } err = g.sessionservice.Delete(ctx, ssReq) - output := "" if err != nil { - output = err.Error() + return fmt.Errorf("g.sessionservice.Delete() failed: %v", err) } - result := models.DeleteSessionResponse{ - Output: output, - } + result := models.DeleteSessionResponse{} err = json.NewEncoder(rw).Encode(result) if err != nil { return fmt.Errorf("json.NewEncoder failed: %v", err) diff --git a/server/agentengine/controllers/method/stream_query.go b/server/agentengine/controllers/method/stream_query.go index 0337adc8c..0196c5b57 100644 --- a/server/agentengine/controllers/method/stream_query.go +++ b/server/agentengine/controllers/method/stream_query.go @@ -60,11 +60,27 @@ func (s *streamQueryHandler) Handle(ctx context.Context, rw http.ResponseWriter, func (s *streamQueryHandler) streamJSONL(ctx context.Context, rw http.ResponseWriter, payload []byte) error { var req models.StreamQueryRequest + // try to unmarshal models.StreamQueryRequest first err := json.Unmarshal(payload, &req) if err != nil { - err = fmt.Errorf("json.Unmarshal() failed: %v", err) - log.Print(err.Error()) - return err + // try to unmarshal models.StreamQueryTextRequest + var reqText models.StreamQueryTextRequest + errText := json.Unmarshal(payload, &reqText) + if errText != nil { + // cannot unmarshall to models.StreamQueryRequest and models.StreamQueryTextRequest + err = fmt.Errorf("json.Unmarshal() failed both for models.StreamQueryRequest (%v) and models.StreamQueryTextRequest (%v)", err, errText) + log.Print(err.Error()) + return err + } + // got text, create a full content based on that text + req = models.StreamQueryRequest{ + ClassMethod: reqText.ClassMethod, + Input: models.StreamQueryInput{ + UserID: reqText.Input.UserID, + SessionID: reqText.Input.SessionID, + Message: *genai.NewContentFromText(reqText.Input.Message, genai.RoleUser), + }, + } } events, err := s.run(ctx, &req, &req.Input.Message, s.config) @@ -97,8 +113,7 @@ func (s *streamQueryHandler) streamJSONL(ctx context.Context, rw http.ResponseWr continue } - chunk := *event - err = helper.EmitJSON(rw, chunk) + err = helper.EmitJSON(rw, *event) if err != nil { e := fmt.Errorf("helper.EmitJSON() failed: %w", err) log.Print(e.Error()) @@ -135,8 +150,15 @@ func (s *streamQueryHandler) Metadata() (*structpb.Struct, error) { "type": "string", }, "message": map[string]any{ - "additionalProperties": true, - "type": "object", + "anyOf": []any{ + map[string]any{ + "type": "string", + }, + map[string]any{ + "additionalProperties": true, + "type": "object", + }, + }, }, }, "required": []any{ @@ -169,11 +191,13 @@ func (s *streamQueryHandler) run(ctx context.Context, req *models.StreamQueryReq rootAgent := config.AgentLoader.RootAgent() r, err := runner.New(runner.Config{ - AppName: s.agentEngineID, - Agent: rootAgent, - SessionService: config.SessionService, - ArtifactService: config.ArtifactService, - PluginConfig: config.PluginConfig, + AppName: s.agentEngineID, + Agent: rootAgent, + SessionService: config.SessionService, + MemoryService: config.MemoryService, + ArtifactService: config.ArtifactService, + PluginConfig: config.PluginConfig, + AutoCreateSession: true, }) if err != nil { return nil, fmt.Errorf("failed to create runner: %v", err) diff --git a/server/agentengine/controllers/method/stream_query_test.go b/server/agentengine/controllers/method/stream_query_test.go new file mode 100644 index 000000000..4a7f35e97 --- /dev/null +++ b/server/agentengine/controllers/method/stream_query_test.go @@ -0,0 +1,165 @@ +// Copyright 2026 Google LLC +// +// Licensed under the Apache License, Version 2.0 (the "License"); +// you may not use this file except in compliance with the License. +// You may obtain a copy of the License at +// +// http://www.apache.org/licenses/LICENSE-2.0 +// +// Unless required by applicable law or agreed to in writing, software +// distributed under the License is distributed on an "AS IS" BASIS, +// WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. +// See the License for the specific language governing permissions and +// limitations under the License. + +package method + +import ( + "encoding/json" + "net/http" + "strconv" + "strings" + "testing" + + "google.golang.org/genai" + + "google.golang.org/adk/agent" + "google.golang.org/adk/agent/llmagent" + "google.golang.org/adk/cmd/launcher" + "google.golang.org/adk/session" +) + +type simpleEvent struct { + Content *genai.Content `json:"content"` +} + +// TestSimpleText checks whether a simple message as string gives the same result as genai.Content. +func TestSimpleText(t *testing.T) { + agentEngineId := 123 + appName := strconv.Itoa(agentEngineId) + userID := "u" + + // agent invokes BeforeAgent callback which returns the content as provided as an answer + a, err := llmagent.New(llmagent.Config{ + Name: "Echo", + BeforeAgentCallbacks: []agent.BeforeAgentCallback{ + func(cc agent.CallbackContext) (*genai.Content, error) { + return cc.UserContent(), nil + }, + }, + }) + if err != nil { + t.Fatalf("failed to create agent: %v", err) + } + + config := &launcher.Config{ + AgentLoader: agent.NewSingleLoader(a), + SessionService: session.InMemoryService(), + } + h := NewStreamQueryHandler(config, appName, "async_stream_query", "") + + ctx := t.Context() + sess, err := config.SessionService.Create(ctx, &session.CreateRequest{AppName: appName, UserID: userID}) + if err != nil { + t.Fatalf("failed to create session: %v", err) + } + + wantContent := genai.NewContentFromText("Say hello", genai.RoleUser) + wantBytes, err := json.Marshal(wantContent) + if err != nil { + t.Fatalf("json.Marshal() failed: %v", err) + } + want := string(wantBytes) + + tests := []struct { + name string + payload string + }{ + { + name: "full content", + payload: `{ +"class_method":"async_stream_query", +"input":{ + "message":{ + "parts":[ + {"text":"Say hello"} + ], + "role":"user" + }, + "session_id":"` + sess.Session.ID() + `", + "user_id":"` + userID + `"}}`, + }, + { + name: "simplified content", + payload: `{ +"class_method":"async_stream_query", +"input":{ + "message":"Say hello", + "session_id":"` + sess.Session.ID() + `", + "user_id":"` + userID + `"}}`, + }, + } + + for _, tt := range tests { + w := newStringWriter() + b := []byte(tt.payload) + err := h.streamJSONL(t.Context(), w, b) + if err != nil { + t.Fatalf("streamJSONL() failed: %v", err) + } + + var ev simpleEvent + p := w.sb.String() + + err = json.Unmarshal([]byte(p), &ev) + if err != nil { + t.Fatalf("json.Unmarshal() failed: %v", err) + } + gotBytes, err := json.Marshal(ev.Content) + if err != nil { + t.Fatalf("json.Marshal() failed: %v", err) + } + got := string(gotBytes) + if got != want { + t.Errorf("streamJSONL() = %v, want %v", got, want) + } + } +} + +// mock writer for http +type stringWriter struct { + sb strings.Builder + h http.Header +} + +// Header implements [http.ResponseWriter]. +func (s *stringWriter) Header() http.Header { + return s.h +} + +// WriteHeader implements [http.ResponseWriter]. +func (s *stringWriter) WriteHeader(statusCode int) { + s.h = http.Header{"Status": []string{http.StatusText(statusCode)}} +} + +// Write implements [http.ResponseWriter]. +func (s *stringWriter) Write(p []byte) (n int, err error) { + return s.sb.Write(p) +} + +// Flush implements [http.Flusher] +func (s *stringWriter) Flush() { + // do nothing +} + +var ( + _ http.ResponseWriter = (*stringWriter)(nil) + _ http.Flusher = (*stringWriter)(nil) +) + +func newStringWriter() *stringWriter { + return &stringWriter{ + sb: strings.Builder{}, + h: http.Header{}, + } +} diff --git a/server/agentengine/controllers/method/streaming_agent_run_with_events.go b/server/agentengine/controllers/method/streaming_agent_run_with_events.go new file mode 100644 index 000000000..e9cfad7dc --- /dev/null +++ b/server/agentengine/controllers/method/streaming_agent_run_with_events.go @@ -0,0 +1,251 @@ +// Copyright 2026 Google LLC +// +// Licensed under the Apache License, Version 2.0 (the "License"); +// you may not use this file except in compliance with the License. +// You may obtain a copy of the License at +// +// http://www.apache.org/licenses/LICENSE-2.0 +// +// Unless required by applicable law or agreed to in writing, software +// distributed under the License is distributed on an "AS IS" BASIS, +// WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. +// See the License for the specific language governing permissions and +// limitations under the License. + +package method + +import ( + "context" + "encoding/json" + "fmt" + "iter" + "log" + "net/http" + + "google.golang.org/genai" + "google.golang.org/protobuf/types/known/structpb" + + "google.golang.org/adk/agent" + "google.golang.org/adk/cmd/launcher" + "google.golang.org/adk/runner" + "google.golang.org/adk/server/agentengine/internal/helper" + "google.golang.org/adk/server/agentengine/internal/models" + "google.golang.org/adk/session" +) + +type streamingAgentRunWithEventsHandler struct { + config *launcher.Config + methodName string + apiMode string + agentEngineID string +} + +// NewStreamingAgentRunWithEventsHandler creates a new handler for streaming_agent_run_with_events. +func NewStreamingAgentRunWithEventsHandler(config *launcher.Config, agentEngineID, methodName, apiMode string) *streamingAgentRunWithEventsHandler { + return &streamingAgentRunWithEventsHandler{config: config, agentEngineID: agentEngineID, methodName: methodName, apiMode: apiMode} +} + +// Handle generates stream of json-encoded responses based on the payload. Errors are also emitted as errors. +func (s *streamingAgentRunWithEventsHandler) Handle(ctx context.Context, rw http.ResponseWriter, payload []byte) error { + streamErr := s.streamJSONL(ctx, rw, payload) + // streamJSONL will return error only before streaming. In that case we can handle it with HTTP Status, which is done in upstream. + if streamErr != nil { + err := fmt.Errorf("s.streamJSONL() failed: %w", streamErr) + return err + } + return nil +} + +// streamJSONL streams a single line for each event or error. +func (s *streamingAgentRunWithEventsHandler) streamJSONL(ctx context.Context, rw http.ResponseWriter, payload []byte) error { + var req models.StreamingAgentRunWithEventsRequest + if err := json.Unmarshal(payload, &req); err != nil { + err = fmt.Errorf("json.Unmarshal() failed for models.StreamingAgentRunWithEventsRequest: %w", err) + log.Print(err.Error()) + return err + } + + runReq, requestedSessionID, err := decodeStreamingAgentRunWithEventsRequest(&req) + if err != nil { + err = fmt.Errorf("decodeStreamingAgentRunWithEventsRequest() failed: %w", err) + log.Print(err.Error()) + return err + } + if err := s.ensureBackendSession(ctx, runReq, requestedSessionID); err != nil { + err = fmt.Errorf("s.ensureBackendSession() failed: %w", err) + log.Print(err.Error()) + return err + } + + events, err := s.run(ctx, runReq, &runReq.Message, s.config) + if err != nil { + err = fmt.Errorf("s.run() failed: %w", err) + log.Print(err.Error()) + return err + } + + rw.Header().Set("Content-Type", "application/json") + rw.Header().Set("Cache-Control", "no-cache") + rw.Header().Set("Connection", "keep-alive") + // from this moment on we must not return error. Instead, it should be handled by using helper.EmitJSONError + + for event, err := range events { + log.Printf("Processing event: %+v err: %+v\n", event, err) + if err != nil { + log.Printf("error in events: %v\n", err) + e := helper.EmitJSONError(rw, err) + if e != nil { + e = fmt.Errorf("helper.EmitJSONError() failed: %w", e) + log.Print(e.Error()) + } + break + } + if event == nil { + continue + } + if event.LLMResponse.Content == nil { + continue + } + + err = helper.EmitJSON(rw, models.StreamingAgentRunWithEventsResponse{ + Events: []*session.Event{event}, + SessionID: runReq.SessionID, + }) + if err != nil { + e := fmt.Errorf("helper.EmitJSON() failed: %w", err) + log.Print(e.Error()) + e = helper.EmitJSONError(rw, e) + if e != nil { + e = fmt.Errorf("helper.EmitJSONError() failed: %w", e) + log.Print(e.Error()) + } + break + } + } + return nil +} + +// decodeStreamingAgentRunWithEventsRequest decodes input.request_json and returns +// the caller-requested session_id before backend ADK session normalization. +func decodeStreamingAgentRunWithEventsRequest(req *models.StreamingAgentRunWithEventsRequest) (*models.StreamingAgentRunWithEventsRunRequest, string, error) { + var runReq models.StreamingAgentRunWithEventsRunRequest + if err := json.Unmarshal([]byte(req.Input.RequestJSON), &runReq); err != nil { + return nil, "", fmt.Errorf("json.Unmarshal(input.request_json) failed: %w", err) + } + return &runReq, runReq.SessionID, nil +} + +// ensureBackendSession normalizes Gemini Enterprise requests so +// the runner always receives a backend ADK session ID. +// +// On the first turn, the embedded session_id is usually a Gemini Enterprise / +// Discovery Engine session resource: +// +// projects/{project}/locations/global/collections/default_collection/engines/{engine}/sessions/{session} +// +// VertexAISessionService cannot use that resource name as a caller-provided +// backend session ID. This method treats the incoming session_id as either a +// backend ADK session ID returned by a previous response, or an external +// first-turn resource that needs a newly-created backend session. +func (s *streamingAgentRunWithEventsHandler) ensureBackendSession(ctx context.Context, req *models.StreamingAgentRunWithEventsRunRequest, requestedSessionID string) error { + if requestedSessionID == "" { + return nil + } + if req.UserID == "" { + return fmt.Errorf("user_id is required for backend session handling") + } + if s.config == nil || s.config.SessionService == nil { + return fmt.Errorf("session service is required for backend session handling") + } + + getResp, err := s.config.SessionService.Get(ctx, &session.GetRequest{ + AppName: s.agentEngineID, + UserID: req.UserID, + SessionID: requestedSessionID, + }) + if err == nil && getResp.Session != nil { + req.SessionID = getResp.Session.ID() + return nil + } + + createResp, err := s.config.SessionService.Create(ctx, &session.CreateRequest{ + AppName: s.agentEngineID, + UserID: req.UserID, + }) + if err != nil { + return fmt.Errorf("sessionService.Create() failed: %w", err) + } + + req.SessionID = createResp.Session.ID() + return nil +} + +// Name implements MethodHandler. +func (s *streamingAgentRunWithEventsHandler) Name() string { + return s.methodName +} + +var _ MethodHandler = (*streamingAgentRunWithEventsHandler)(nil) + +// Metadata implements MethodHandler. +func (s *streamingAgentRunWithEventsHandler) Metadata() (*structpb.Struct, error) { + classAsyncMethod, err := structpb.NewStruct(map[string]any{ + "api_mode": s.apiMode, + "name": s.methodName, + "parameters": map[string]any{ + "properties": map[string]any{ + "request_json": map[string]any{ + "type": "string", + }, + }, + "required": []any{ + "request_json", + }, + "type": "object", + }, + "description": `Streams responses asynchronously from the ADK application. + +This method is primarily meant for invocation from Gemini Enterprise. + +Args: + request_json (str): + Required. A JSON-encoded request object with: + - user_id (str): Required. The user ID to run the agent for. + - session_id (str): Optional. The session ID. Gemini Enterprise may + send its Discovery Engine session resource on the first turn; later + turns should use the backend ADK session ID returned in the response. + - message (Content): Required. The user message, using the genai + Content JSON shape, for example: + {"role":"user","parts":[{"text":"Hi"}]} + +`, + }) + if err != nil { + return nil, fmt.Errorf("cannot create %s: %w", s.Name(), err) + } + return classAsyncMethod, nil +} + +func (s *streamingAgentRunWithEventsHandler) run(ctx context.Context, req *models.StreamingAgentRunWithEventsRunRequest, message *genai.Content, config *launcher.Config) (iter.Seq2[*session.Event, error], error) { + rootAgent := config.AgentLoader.RootAgent() + + r, err := runner.New(runner.Config{ + AppName: s.agentEngineID, + Agent: rootAgent, + SessionService: config.SessionService, + ArtifactService: config.ArtifactService, + MemoryService: config.MemoryService, + PluginConfig: config.PluginConfig, + AutoCreateSession: true, + }) + if err != nil { + return nil, fmt.Errorf("failed to create runner: %v", err) + } + + // The path mirrors Python AdkApp.streaming_agent_run_with_events, + // which does not force SSE mode. Sending both partial and final events here + // causes Gemini Enterprise to render duplicate answer text. + return r.Run(ctx, req.UserID, req.SessionID, message, agent.RunConfig{ + StreamingMode: agent.StreamingModeNone, + }), nil +} diff --git a/server/agentengine/controllers/method/streaming_agent_run_with_events_test.go b/server/agentengine/controllers/method/streaming_agent_run_with_events_test.go new file mode 100644 index 000000000..64a0cfb1e --- /dev/null +++ b/server/agentengine/controllers/method/streaming_agent_run_with_events_test.go @@ -0,0 +1,357 @@ +// Copyright 2026 Google LLC +// +// Licensed under the Apache License, Version 2.0 (the "License"); +// you may not use this file except in compliance with the License. +// You may obtain a copy of the License at +// +// http://www.apache.org/licenses/LICENSE-2.0 +// +// Unless required by applicable law or agreed to in writing, software +// distributed under the License is distributed on an "AS IS" BASIS, +// WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. +// See the License for the specific language governing permissions and +// limitations under the License. + +package method + +import ( + "context" + "encoding/json" + "iter" + "testing" + + "github.com/google/go-cmp/cmp" + "github.com/google/go-cmp/cmp/cmpopts" + "google.golang.org/genai" + + "google.golang.org/adk/agent" + "google.golang.org/adk/agent/llmagent" + "google.golang.org/adk/cmd/launcher" + "google.golang.org/adk/model" + "google.golang.org/adk/server/agentengine/internal/models" + "google.golang.org/adk/session" +) + +type agentSpaceStreamResponse struct { + Events []simpleEvent `json:"events"` + SessionID string `json:"session_id"` +} + +type streamAwareLLM struct{} + +func (streamAwareLLM) Name() string { + return "stream-aware-llm" +} + +func (streamAwareLLM) GenerateContent(ctx context.Context, req *model.LLMRequest, stream bool) iter.Seq2[*model.LLMResponse, error] { + return func(yield func(*model.LLMResponse, error) bool) { + if stream { + if !yield(&model.LLMResponse{ + Content: genai.NewContentFromText("partial response", genai.RoleModel), + Partial: true, + }, nil) { + return + } + } + yield(&model.LLMResponse{ + Content: genai.NewContentFromText("final response", genai.RoleModel), + }, nil) + } +} + +func TestDecodeStreamingAgentRunWithEventsRequest(t *testing.T) { + payload := []byte(`{ + "class_method": "streaming_agent_run_with_events", + "input": { + "request_json": "{\"message\":{\"role\":\"user\",\"parts\":[{\"text\":\"Hi\"}]},\"session_id\":\"projects/111111111111/locations/global/collections/default_collection/engines/test-engine/sessions/12345678901234567890\",\"user_id\":\"test-user@example.com\"}" + } + }`) + var req models.StreamingAgentRunWithEventsRequest + if err := json.Unmarshal(payload, &req); err != nil { + t.Fatalf("json.Unmarshal() failed: %v", err) + } + + got, requestedSessionID, err := decodeStreamingAgentRunWithEventsRequest(&req) + if err != nil { + t.Fatalf("decodeStreamingAgentRunWithEventsRequest() failed: %v", err) + } + + want := &models.StreamingAgentRunWithEventsRunRequest{ + UserID: "test-user@example.com", + SessionID: "projects/111111111111/locations/global/collections/default_collection/engines/test-engine/sessions/12345678901234567890", + Message: genai.Content{ + Role: "user", + Parts: []*genai.Part{{Text: "Hi"}}, + }, + } + + if diff := cmp.Diff(want, got); diff != "" { + t.Errorf("decodeStreamingAgentRunWithEventsRequest() mismatch (-want +got):\n%s", diff) + } + if requestedSessionID != want.SessionID { + t.Errorf("requestedSessionID = %q, want %q", requestedSessionID, want.SessionID) + } +} + +func TestEnsureBackendSession_CreateBackendSession(t *testing.T) { + sessionService := session.InMemoryService() + handler := NewStreamingAgentRunWithEventsHandler(&launcher.Config{SessionService: sessionService}, "app", "streaming_agent_run_with_events", "async_stream") + req := &models.StreamingAgentRunWithEventsRunRequest{ + UserID: "test-user@example.com", + SessionID: "projects/111111111111/locations/global/collections/default_collection/engines/test-engine/sessions/12345678901234567890", + } + requestedSessionID := req.SessionID + + if err := handler.ensureBackendSession(t.Context(), req, requestedSessionID); err != nil { + t.Fatalf("ensureBackendSession() failed: %v", err) + } + if req.SessionID == "" || req.SessionID == requestedSessionID { + t.Fatalf("SessionID = %q, want generated backend session ID", req.SessionID) + } + + got, err := sessionService.Get(t.Context(), &session.GetRequest{ + AppName: "app", + UserID: "test-user@example.com", + SessionID: req.SessionID, + }) + if err != nil { + t.Fatalf("Get() failed: %v", err) + } + if got.Session.ID() != req.SessionID { + t.Errorf("stored SessionID = %q, want %q", got.Session.ID(), req.SessionID) + } +} + +func TestEnsureBackendSession_ReuseReturnedBackendSession(t *testing.T) { + sessionService := session.InMemoryService() + created, err := sessionService.Create(t.Context(), &session.CreateRequest{ + AppName: "app", + UserID: "jan@example.com", + }) + if err != nil { + t.Fatalf("Create() failed: %v", err) + } + + handler := NewStreamingAgentRunWithEventsHandler(&launcher.Config{SessionService: sessionService}, "app", "streaming_agent_run_with_events", "async_stream") + req := &models.StreamingAgentRunWithEventsRunRequest{ + UserID: "jan@example.com", + SessionID: created.Session.ID(), + } + + if err := handler.ensureBackendSession(t.Context(), req, req.SessionID); err != nil { + t.Fatalf("ensureBackendSession() failed: %v", err) + } + if req.SessionID != created.Session.ID() { + t.Errorf("SessionID = %q, want existing backend session %q", req.SessionID, created.Session.ID()) + } +} + +func TestStreamJSONL_AgentSpaceResponseEnvelope(t *testing.T) { + const ( + appName = "app" + userID = "test-user@example.com" + externalSessionID = "projects/111111111111/locations/global/collections/default_collection/engines/test-engine/sessions/12345678901234567890" + ) + + a, err := llmagent.New(llmagent.Config{ + Name: "Echo", + BeforeAgentCallbacks: []agent.BeforeAgentCallback{ + func(cc agent.CallbackContext) (*genai.Content, error) { + return cc.UserContent(), nil + }, + }, + }) + if err != nil { + t.Fatalf("failed to create agent: %v", err) + } + + config := &launcher.Config{ + AgentLoader: agent.NewSingleLoader(a), + SessionService: session.InMemoryService(), + } + h := NewStreamingAgentRunWithEventsHandler(config, appName, "streaming_agent_run_with_events", "async_stream") + + requestJSON := `{"message":{"role":"user","parts":[{"text":"Please"}]},"session_id":"` + externalSessionID + `","user_id":"` + userID + `"}` + payload, err := json.Marshal(models.StreamingAgentRunWithEventsRequest{ + ClassMethod: "streaming_agent_run_with_events", + Input: models.StreamingAgentRunWithEventsInput{ + RequestJSON: requestJSON, + }, + }) + if err != nil { + t.Fatalf("json.Marshal() failed: %v", err) + } + + w := newStringWriter() + if err := h.streamJSONL(t.Context(), w, payload); err != nil { + t.Fatalf("streamJSONL() failed: %v", err) + } + + var got agentSpaceStreamResponse + if err := json.Unmarshal([]byte(w.sb.String()), &got); err != nil { + t.Fatalf("json.Unmarshal() failed: %v", err) + } + if got.SessionID == "" || got.SessionID == externalSessionID { + t.Fatalf("SessionID = %q, want generated backend session ID", got.SessionID) + } + if len(got.Events) != 1 { + t.Fatalf("len(Events) = %d, want 1", len(got.Events)) + } + + wantContent := genai.NewContentFromText("Please", genai.RoleUser) + if diff := cmp.Diff(wantContent, got.Events[0].Content); diff != "" { + t.Errorf("event content mismatch (-want +got):\n%s", diff) + } +} + +func TestStreamJSONL_AgentSpaceAcceptsReturnedBackendSessionID(t *testing.T) { + const ( + appName = "app" + userID = "test-user@example.com" + externalSessionID = "projects/111111111111/locations/global/collections/default_collection/engines/test-engine/sessions/12345678901234567890" + ) + + a, err := llmagent.New(llmagent.Config{ + Name: "Echo", + BeforeAgentCallbacks: []agent.BeforeAgentCallback{ + func(cc agent.CallbackContext) (*genai.Content, error) { + return cc.UserContent(), nil + }, + }, + }) + if err != nil { + t.Fatalf("failed to create agent: %v", err) + } + + config := &launcher.Config{ + AgentLoader: agent.NewSingleLoader(a), + SessionService: session.InMemoryService(), + } + h := NewStreamingAgentRunWithEventsHandler(config, appName, "streaming_agent_run_with_events", "async_stream") + + run := func(message, sessionID string) agentSpaceStreamResponse { + t.Helper() + requestJSON := `{"message":{"role":"user","parts":[{"text":"` + message + `"}]},"session_id":"` + sessionID + `","user_id":"` + userID + `"}` + payload, err := json.Marshal(models.StreamingAgentRunWithEventsRequest{ + ClassMethod: "streaming_agent_run_with_events", + Input: models.StreamingAgentRunWithEventsInput{ + RequestJSON: requestJSON, + }, + }) + if err != nil { + t.Fatalf("json.Marshal() failed: %v", err) + } + + w := newStringWriter() + if err := h.streamJSONL(t.Context(), w, payload); err != nil { + t.Fatalf("streamJSONL() failed: %v", err) + } + + var got agentSpaceStreamResponse + if err := json.Unmarshal([]byte(w.sb.String()), &got); err != nil { + t.Fatalf("json.Unmarshal() failed: %v", err) + } + return got + } + + first := run("Hi", externalSessionID) + second := run("Again", first.SessionID) + + if first.SessionID == "" || first.SessionID == externalSessionID { + t.Fatalf("first SessionID = %q, want generated backend session ID", first.SessionID) + } + if second.SessionID != first.SessionID { + t.Fatalf("second SessionID = %q, want returned backend session ID %q", second.SessionID, first.SessionID) + } + + list, err := config.SessionService.List(t.Context(), &session.ListRequest{ + AppName: appName, + UserID: userID, + }) + if err != nil { + t.Fatalf("List() failed: %v", err) + } + if len(list.Sessions) != 1 { + t.Fatalf("len(Sessions) = %d, want 1", len(list.Sessions)) + } +} + +func TestStreamJSONL_AgentSpaceUsesNonStreamingMode(t *testing.T) { + const ( + appName = "app" + userID = "test-user@example.com" + externalSessionID = "projects/111111111111/locations/global/collections/default_collection/engines/test-engine/sessions/12345678901234567890" + ) + + a, err := llmagent.New(llmagent.Config{ + Name: "StreamAware", + Model: streamAwareLLM{}, + }) + if err != nil { + t.Fatalf("failed to create agent: %v", err) + } + + config := &launcher.Config{ + AgentLoader: agent.NewSingleLoader(a), + SessionService: session.InMemoryService(), + } + h := NewStreamingAgentRunWithEventsHandler(config, appName, "streaming_agent_run_with_events", "async_stream") + + requestJSON := `{"message":{"role":"user","parts":[{"text":"What is your capabilities"}]},"session_id":"` + externalSessionID + `","user_id":"` + userID + `"}` + payload, err := json.Marshal(models.StreamingAgentRunWithEventsRequest{ + ClassMethod: "streaming_agent_run_with_events", + Input: models.StreamingAgentRunWithEventsInput{ + RequestJSON: requestJSON, + }, + }) + if err != nil { + t.Fatalf("json.Marshal() failed: %v", err) + } + + w := newStringWriter() + if err := h.streamJSONL(t.Context(), w, payload); err != nil { + t.Fatalf("streamJSONL() failed: %v", err) + } + + var got agentSpaceStreamResponse + if err := json.Unmarshal([]byte(w.sb.String()), &got); err != nil { + t.Fatalf("json.Unmarshal() failed: %v", err) + } + if len(got.Events) != 1 { + t.Fatalf("len(Events) = %d, want 1", len(got.Events)) + } + + wantContent := genai.NewContentFromText("final response", genai.RoleModel) + if diff := cmp.Diff(wantContent, got.Events[0].Content); diff != "" { + t.Errorf("event content mismatch (-want +got):\n%s", diff) + } +} + +func TestStreamingAgentRunWithEventsHandlerMetadata(t *testing.T) { + handler := NewStreamingAgentRunWithEventsHandler(nil, "", "streaming_agent_run_with_events", "async_stream") + + got, err := handler.Metadata() + if err != nil { + t.Fatalf("Metadata() failed: %v", err) + } + + want := map[string]any{ + "api_mode": "async_stream", + "name": "streaming_agent_run_with_events", + "parameters": map[string]any{ + "properties": map[string]any{ + "request_json": map[string]any{ + "type": "string", + }, + }, + "required": []any{"request_json"}, + "type": "object", + }, + } + + if diff := cmp.Diff(want, got.AsMap(), cmpopts.IgnoreMapEntries(func(k string, _ any) bool { + return k == "description" + })); diff != "" { + t.Errorf("Metadata() mismatch (-want +got):\n%s", diff) + } +} diff --git a/server/agentengine/handler.go b/server/agentengine/handler.go index 74aeb063c..be13e0ba2 100644 --- a/server/agentengine/handler.go +++ b/server/agentengine/handler.go @@ -93,6 +93,7 @@ func listNonStreamHandlers(config *launcher.Config, agentEngineID string) []meth func listStreamHandlers(config *launcher.Config, agentEngineID string) []method.MethodHandler { return []method.MethodHandler{ method.NewStreamQueryHandler(config, agentEngineID, "async_stream_query", "async_stream"), + method.NewStreamingAgentRunWithEventsHandler(config, agentEngineID, "streaming_agent_run_with_events", "async_stream"), } } diff --git a/server/agentengine/internal/helper/encode.go b/server/agentengine/internal/helper/encode.go index ab99078f5..8e9eb905d 100644 --- a/server/agentengine/internal/helper/encode.go +++ b/server/agentengine/internal/helper/encode.go @@ -145,7 +145,7 @@ func convertSnake(path, indent string, o any) (any, error) { return map[string]any{}, nil } return res, nil - case reflect.Ptr: + case reflect.Pointer: if v.IsNil() { return nil, nil } diff --git a/server/agentengine/internal/models/session.go b/server/agentengine/internal/models/session.go index e84c6a9b8..b2a707370 100644 --- a/server/agentengine/internal/models/session.go +++ b/server/agentengine/internal/models/session.go @@ -26,9 +26,8 @@ type CreateSessionRequest struct { // CreateSessionInput contains input parameters type CreateSessionInput struct { - UserID string `json:"user_id"` - SessionID string `json:"session_id,omitempty"` - State map[string]any `json:"state,omitempty"` + UserID string `json:"user_id"` + State map[string]any `json:"state,omitempty"` } // CreateSessionResponse contains response diff --git a/server/agentengine/internal/models/stream_query.go b/server/agentengine/internal/models/stream_query.go index dcf7fec07..628584f86 100644 --- a/server/agentengine/internal/models/stream_query.go +++ b/server/agentengine/internal/models/stream_query.go @@ -20,7 +20,7 @@ import ( "google.golang.org/adk/session" ) -// StreamQueryRequest is a struct representing JSON-encoded payload to async_stream_query method with dedicated Input. +// StreamQueryRequest is a struct representing JSON-encoded payload to async_stream_query method with dedicated Input with full genai.Content. type StreamQueryRequest struct { ClassMethod string `json:"class_method"` Input StreamQueryInput `json:"input"` @@ -33,6 +33,19 @@ type StreamQueryInput struct { Message genai.Content `json:"message"` } +// StreamQueryTextRequest is a struct representing JSON-encoded payload to async_stream_query method with dedicated Input with simple text as the content. +type StreamQueryTextRequest struct { + ClassMethod string `json:"class_method"` + Input StreamQueryTextInput `json:"input"` +} + +// StreamQueryTextInput is the actual Input for async_stream_query method. +type StreamQueryTextInput struct { + UserID string `json:"user_id"` + SessionID string `json:"session_id"` + Message string `json:"message"` +} + // StreamQueryResponse defines the content of event data for async_stream_query method. // It is returned as one line with JSON-encoded StreamQuerySSEResponse // Please mind that errors are also returned by this method diff --git a/server/agentengine/internal/models/streaming_agent_run_with_events.go b/server/agentengine/internal/models/streaming_agent_run_with_events.go new file mode 100644 index 000000000..c5e89331a --- /dev/null +++ b/server/agentengine/internal/models/streaming_agent_run_with_events.go @@ -0,0 +1,50 @@ +// Copyright 2026 Google LLC +// +// Licensed under the Apache License, Version 2.0 (the "License"); +// you may not use this file except in compliance with the License. +// You may obtain a copy of the License at +// +// http://www.apache.org/licenses/LICENSE-2.0 +// +// Unless required by applicable law or agreed to in writing, software +// distributed under the License is distributed on an "AS IS" BASIS, +// WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. +// See the License for the specific language governing permissions and +// limitations under the License. + +package models + +import ( + "google.golang.org/genai" + + "google.golang.org/adk/session" +) + +// StreamingAgentRunWithEventsRequest is the JSON-encoded payload for the +// streaming_agent_run_with_events method. +type StreamingAgentRunWithEventsRequest struct { + ClassMethod string `json:"class_method"` + Input StreamingAgentRunWithEventsInput `json:"input"` +} + +// StreamingAgentRunWithEventsInput wraps the actual request payload as JSON. +// RequestJSON is a JSON-encoded StreamingAgentRunWithEventsRunRequest. +type StreamingAgentRunWithEventsInput struct { + RequestJSON string `json:"request_json"` +} + +// StreamingAgentRunWithEventsRunRequest is the request decoded from +// StreamingAgentRunWithEventsInput.RequestJSON. +type StreamingAgentRunWithEventsRunRequest struct { + UserID string `json:"user_id"` + SessionID string `json:"session_id"` + Message genai.Content `json:"message"` +} + +// StreamingAgentRunWithEventsResponse is the response envelope expected by +// Gemini Enterprise for streaming_agent_run_with_events. +type StreamingAgentRunWithEventsResponse struct { + Events []*session.Event `json:"events,omitempty"` + Artifacts []any `json:"artifacts,omitempty"` + SessionID string `json:"session_id,omitempty"` +} diff --git a/session/database/service.go b/session/database/service.go index ec64808f1..b4fc80291 100644 --- a/session/database/service.go +++ b/session/database/service.go @@ -22,9 +22,9 @@ import ( "strings" "time" - "github.com/google/uuid" "gorm.io/gorm" + "google.golang.org/adk/platform" "google.golang.org/adk/session" ) @@ -75,7 +75,7 @@ func (s *databaseService) Create(ctx context.Context, req *session.CreateRequest sessionID := req.SessionID if sessionID == "" { - sessionID = uuid.NewString() + sessionID = platform.NewUUID(ctx) } stateMap := req.State @@ -87,9 +87,9 @@ func (s *databaseService) Create(ctx context.Context, req *session.CreateRequest userID: req.UserID, sessionID: sessionID, state: stateMap, - updatedAt: time.Now(), + updatedAt: platform.Now(ctx), } - createdSession, err := createStorageSession(val) + createdSession, err := createStorageSession(ctx, val) if err != nil { return nil, err } diff --git a/session/database/service_test.go b/session/database/service_test.go index 90de7adcb..59e21efab 100644 --- a/session/database/service_test.go +++ b/session/database/service_test.go @@ -16,21 +16,41 @@ package database import ( "testing" + "time" "github.com/glebarez/sqlite" "gorm.io/gorm" + "google.golang.org/adk/platform" "google.golang.org/adk/session" - "google.golang.org/adk/session/session_test" + "google.golang.org/adk/session/sessiontestsuite" ) func Test_databaseService(t *testing.T) { - opts := session_test.SuiteOptions{SupportsUserProvidedSessionID: true} - session_test.RunServiceTests(t, opts, func(t *testing.T) session.Service { + opts := sessiontestsuite.SuiteOptions{SupportsUserProvidedSessionID: true} + sessiontestsuite.RunServiceTests(t, opts, func(t *testing.T) session.Service { return emptyService(t) }) } +func Test_databaseService_CreateUsesProviders(t *testing.T) { + fixedTime := time.Date(2024, time.January, 2, 3, 4, 5, 0, time.UTC) + ctx := platform.WithTimeProvider(t.Context(), func() time.Time { return fixedTime }) + ctx = platform.WithUUIDProvider(ctx, func() string { return "fixed-session-id" }) + + s := emptyService(t) + resp, err := s.Create(ctx, &session.CreateRequest{AppName: "app", UserID: "user"}) + if err != nil { + t.Fatalf("Create() error = %v", err) + } + if got := resp.Session.ID(); got != "fixed-session-id" { + t.Errorf("autogenerated session ID = %q, want provider value %q", got, "fixed-session-id") + } + if got := resp.Session.LastUpdateTime(); !got.Equal(fixedTime) { + t.Errorf("LastUpdateTime() = %v, want provider value %v", got, fixedTime) + } +} + func emptyService(t *testing.T) *databaseService { t.Helper() gormConfig := &gorm.Config{ diff --git a/session/database/storage_session.go b/session/database/storage_session.go index b773227f9..7d00c917b 100644 --- a/session/database/storage_session.go +++ b/session/database/storage_session.go @@ -15,6 +15,7 @@ package database import ( + "context" "encoding/json" "fmt" "time" @@ -22,6 +23,7 @@ import ( "google.golang.org/genai" "google.golang.org/adk/model" + "google.golang.org/adk/platform" "google.golang.org/adk/session" ) @@ -44,14 +46,15 @@ func (storageSession) TableName() string { } // Helper to map from internal struct to GORM struct -func createStorageSession(s *localSession) (*storageSession, error) { +func createStorageSession(ctx context.Context, s *localSession) (*storageSession, error) { + now := platform.Now(ctx) return &storageSession{ UserID: s.userID, AppName: s.appName, ID: s.sessionID, State: s.state, - CreateTime: time.Now(), - UpdateTime: time.Now(), + CreateTime: now, + UpdateTime: now, }, nil } diff --git a/session/event_test.go b/session/event_test.go new file mode 100644 index 000000000..b2dcbe980 --- /dev/null +++ b/session/event_test.go @@ -0,0 +1,99 @@ +// Copyright 2026 Google LLC +// +// Licensed under the Apache License, Version 2.0 (the "License"); +// you may not use this file except in compliance with the License. +// You may obtain a copy of the License at +// +// http://www.apache.org/licenses/LICENSE-2.0 +// +// Unless required by applicable law or agreed to in writing, software +// distributed under the License is distributed on an "AS IS" BASIS, +// WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. +// See the License for the specific language governing permissions and +// limitations under the License. + +package session_test + +import ( + "context" + "testing" + "time" + + "github.com/google/uuid" + + "google.golang.org/adk/platform" + "google.golang.org/adk/session" +) + +// TestNewEventDefaults covers the deprecated NewEvent, which must keep its +// original signature and use the wall clock and a random UUID. +func TestNewEventDefaults(t *testing.T) { + before := time.Now() + ev := session.NewEvent("inv-1") + after := time.Now() + + if ev.InvocationID != "inv-1" { + t.Errorf("InvocationID = %q, want %q", ev.InvocationID, "inv-1") + } + if _, err := uuid.Parse(ev.ID); err != nil { + t.Errorf("ID = %q, not a valid UUID: %v", ev.ID, err) + } + if ev.Timestamp.Before(before) || ev.Timestamp.After(after) { + t.Errorf("Timestamp = %v, want within [%v, %v]", ev.Timestamp, before, after) + } +} + +func TestNewEventUsesProviders(t *testing.T) { + fixedTime := time.Date(2024, time.January, 2, 3, 4, 5, 0, time.UTC) + ctx := platform.WithTimeProvider(context.Background(), func() time.Time { return fixedTime }) + ctx = platform.WithUUIDProvider(ctx, func() string { return "fixed-event-id" }) + + ev := session.NewEventWithContext(ctx, "inv-1") + + if ev.ID != "fixed-event-id" { + t.Errorf("ID = %q, want %q", ev.ID, "fixed-event-id") + } + if !ev.Timestamp.Equal(fixedTime) { + t.Errorf("Timestamp = %v, want %v", ev.Timestamp, fixedTime) + } + if ev.InvocationID != "inv-1" { + t.Errorf("InvocationID = %q, want %q", ev.InvocationID, "inv-1") + } +} + +// TestNewEventDeterministicReplay verifies that replaying the same sequence of +// provider values produces identical events. This is the property a workflow +// engine relies on to make event creation replay-safe. +func TestNewEventDeterministicReplay(t *testing.T) { + newCtx := func() context.Context { + var ids int + ctx := platform.WithUUIDProvider(context.Background(), func() string { + ids++ + return "event-" + string(rune('0'+ids)) + }) + var times int + return platform.WithTimeProvider(ctx, func() time.Time { + times++ + return time.Date(2024, time.January, 1, 0, 0, times, 0, time.UTC) + }) + } + + run := func(ctx context.Context) []*session.Event { + return []*session.Event{ + session.NewEventWithContext(ctx, "inv"), + session.NewEventWithContext(ctx, "inv"), + } + } + + first := run(newCtx()) + second := run(newCtx()) + + for i := range first { + if first[i].ID != second[i].ID { + t.Errorf("event %d ID: first run %q, second run %q", i, first[i].ID, second[i].ID) + } + if !first[i].Timestamp.Equal(second[i].Timestamp) { + t.Errorf("event %d Timestamp: first run %v, second run %v", i, first[i].Timestamp, second[i].Timestamp) + } + } +} diff --git a/session/inmemory.go b/session/inmemory.go index 62967e697..24e70949b 100644 --- a/session/inmemory.go +++ b/session/inmemory.go @@ -25,11 +25,11 @@ import ( "sync" "time" - "github.com/google/uuid" "rsc.io/omap" "rsc.io/ordered" "google.golang.org/adk/internal/sessionutils" + "google.golang.org/adk/platform" ) type stateMap map[string]any @@ -50,7 +50,7 @@ func (s *inMemoryService) Create(ctx context.Context, req *CreateRequest) (*Crea sessionID := req.SessionID if sessionID == "" { - sessionID = uuid.NewString() + sessionID = platform.NewUUID(ctx) } key := id{ @@ -74,7 +74,7 @@ func (s *inMemoryService) Create(ctx context.Context, req *CreateRequest) (*Crea val := &session{ id: key, state: state, - updatedAt: time.Now(), + updatedAt: platform.Now(ctx), } s.sessions.Set(encodedKey, val) diff --git a/session/inmemory_test.go b/session/inmemory_test.go index 1d178be89..d98664ca7 100644 --- a/session/inmemory_test.go +++ b/session/inmemory_test.go @@ -21,13 +21,32 @@ import ( "testing" "time" + "google.golang.org/adk/platform" "google.golang.org/adk/session" - "google.golang.org/adk/session/session_test" + "google.golang.org/adk/session/sessiontestsuite" ) +func Test_inMemoryService_CreateUsesProviders(t *testing.T) { + fixedTime := time.Date(2024, time.January, 2, 3, 4, 5, 0, time.UTC) + ctx := platform.WithTimeProvider(t.Context(), func() time.Time { return fixedTime }) + ctx = platform.WithUUIDProvider(ctx, func() string { return "fixed-session-id" }) + + s := session.InMemoryService() + resp, err := s.Create(ctx, &session.CreateRequest{AppName: "app", UserID: "user"}) + if err != nil { + t.Fatalf("Create() error = %v", err) + } + if got := resp.Session.ID(); got != "fixed-session-id" { + t.Errorf("autogenerated session ID = %q, want provider value %q", got, "fixed-session-id") + } + if got := resp.Session.LastUpdateTime(); !got.Equal(fixedTime) { + t.Errorf("LastUpdateTime() = %v, want provider value %v", got, fixedTime) + } +} + func Test_inMemoryService(t *testing.T) { - opts := session_test.SuiteOptions{SupportsUserProvidedSessionID: true} // InMemory supports custom IDs - session_test.RunServiceTests(t, opts, func(t *testing.T) session.Service { + opts := sessiontestsuite.SuiteOptions{SupportsUserProvidedSessionID: true} // InMemory supports custom IDs + sessiontestsuite.RunServiceTests(t, opts, func(t *testing.T) session.Service { return session.InMemoryService() }) } diff --git a/session/live_diagnostics.go b/session/live_diagnostics.go index ed6b018ce..956f3d309 100644 --- a/session/live_diagnostics.go +++ b/session/live_diagnostics.go @@ -17,8 +17,10 @@ package session import "time" // LiveDiagnostics captures computed timing and protocol state for a live -// streaming event. Only populated for events produced by RunLive; nil for -// standard Run() events. +// streaming event. Only populated for events produced by the runner's +// RunLiveQueue; nil for standard Run() events and for events produced by +// the upstream-API RunLive (agent.LiveSession signature), whose engine does +// not compute diagnostics. // // EPHEMERAL: This struct is NOT persisted to storage. It is attached to // yielded events for the caller's use only. The runner strips it before diff --git a/session/session.go b/session/session.go index 9e05bf9a6..7f5935687 100644 --- a/session/session.go +++ b/session/session.go @@ -15,6 +15,7 @@ package session import ( + "context" "errors" "iter" "time" @@ -22,6 +23,7 @@ import ( "github.com/google/uuid" "google.golang.org/adk/model" + "google.golang.org/adk/platform" "google.golang.org/adk/tool/toolconfirmation" ) @@ -116,9 +118,9 @@ type Event struct { // Only valid for function call event. LongRunningToolIDs []string - // LiveDiagnostics is populated only for events from RunLive sessions. - // EPHEMERAL: not persisted to storage. Nil for standard Run() events - // and for events loaded from storage. + // LiveDiagnostics is populated only for events from RunLiveQueue sessions. + // EPHEMERAL: not persisted to storage. Nil for standard Run() events, for + // events from the upstream-API RunLive, and for events loaded from storage. LiveDiagnostics *LiveDiagnostics } @@ -135,6 +137,11 @@ func (e *Event) IsFinalResponse() bool { } // NewEvent creates a new event defining now as the timestamp. +// +// Deprecated: Use [NewEventWithContext] instead so that platform-installed time +// and UUID providers (see [platform.WithTimeProvider] and +// [platform.WithUUIDProvider]) are honored. NewEvent always uses the wall clock +// and a random UUID. func NewEvent(invocationID string) *Event { return &Event{ ID: uuid.NewString(), @@ -144,6 +151,21 @@ func NewEvent(invocationID string) *Event { } } +// NewEventWithContext creates a new event defining now as the timestamp. +// +// The event ID and timestamp are obtained through the platform package, so a +// time or UUID provider installed on ctx (see [platform.WithTimeProvider] and +// [platform.WithUUIDProvider]) controls them. This lets callers such as +// workflow engines produce deterministic, replay-safe events. +func NewEventWithContext(ctx context.Context, invocationID string) *Event { + return &Event{ + ID: platform.NewUUID(ctx), + InvocationID: invocationID, + Timestamp: platform.Now(ctx), + Actions: EventActions{StateDelta: make(map[string]any), ArtifactDelta: make(map[string]int64)}, + } +} + // EventActions represent the actions attached to an event. type EventActions struct { // Set by agent.Context implementation. diff --git a/session/sessiontestsuite/service_suite.go b/session/sessiontestsuite/service_suite.go new file mode 100644 index 000000000..2639e194b --- /dev/null +++ b/session/sessiontestsuite/service_suite.go @@ -0,0 +1,752 @@ +// Copyright 2026 Google LLC +// +// Licensed under the Apache License, Version 2.0 (the "License"); +// you may not use this file except in compliance with the License. +// You may obtain a copy of the License at +// +// http://www.apache.org/licenses/LICENSE-2.0 +// +// Unless required by applicable law or agreed to in writing, software +// distributed under the License is distributed on an "AS IS" BASIS, +// WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. +// See the License for the specific language governing permissions and +// limitations under the License. + +package sessiontestsuite + +import ( + "strconv" + "testing" + "time" + + "github.com/google/go-cmp/cmp" + "github.com/google/go-cmp/cmp/cmpopts" + "google.golang.org/genai" + "google.golang.org/grpc/codes" + "google.golang.org/grpc/status" + + "google.golang.org/adk/model" + "google.golang.org/adk/session" +) + +// ExpectedSession represents a snapshot of a session's public state for test comparisons. +type ExpectedSession struct { + AppName string + UserID string + SessionID string + State map[string]any + Events []*session.Event +} + +// Snapshot extracts values from ANY session.Session implementation. +func Snapshot(s session.Session) ExpectedSession { + if s == nil { + return ExpectedSession{} + } + + state := make(map[string]any) + if s.State() != nil { + for k, v := range s.State().All() { + state[k] = v + } + } + + var events []*session.Event + if s.Events() != nil { + for e := range s.Events().All() { + events = append(events, e) + } + } + + return ExpectedSession{ + AppName: s.AppName(), + UserID: s.UserID(), + SessionID: s.ID(), + State: state, + Events: events, + } +} + +// SuiteOptions holds configuration for adaptive test runs. +type SuiteOptions struct { + SupportsUserProvidedSessionID bool + ProvidesServerAssignedEventID bool + AppName string +} + +// RunServiceTests runs a battery of standard tests against a Session.Service. +func RunServiceTests(t *testing.T, opts SuiteOptions, setup func(t *testing.T) session.Service) { + testAppName := "testApp" + if opts.AppName != "" { + testAppName = opts.AppName + } + t.Run("Create", func(t *testing.T) { + t.Run("full_key", func(t *testing.T) { + if !opts.SupportsUserProvidedSessionID { + t.Skip("Skipping full key test: requires user provided session ID support") + } + s := setup(t) + req := &session.CreateRequest{ + AppName: testAppName, + UserID: "testUserID", + SessionID: "test-session-id", + State: map[string]any{ + "k": float64(5), + }, + } + + got, err := s.Create(t.Context(), req) + if opts.SupportsUserProvidedSessionID { + if err != nil { + t.Fatalf("Create() error = %v, wantErr %v", err, false) + } + if got.Session.AppName() != req.AppName { + t.Errorf("AppName got: %v, want: %v", got.Session.AppName(), req.AppName) + } + if got.Session.UserID() != req.UserID { + t.Errorf("UserID got: %v, want: %v", got.Session.UserID(), req.UserID) + } + if got.Session.ID() != req.SessionID { + t.Errorf("SessionID got: %v, want: %v", got.Session.ID(), req.SessionID) + } + } else { + if err == nil { + t.Fatalf("Expected error for user-provided SessionID") + } + } + }) + + t.Run("generated_session_id", func(t *testing.T) { + s := setup(t) + req := &session.CreateRequest{ + AppName: testAppName, + UserID: "testUserID", + State: map[string]any{"k": float64(5)}, + } + + got, err := s.Create(t.Context(), req) + if err != nil { + t.Fatalf("Create() error = %v", err) + } + if got.Session.ID() == "" { + t.Errorf("Expected generated SessionID, got empty") + } + }) + + t.Run("when_already_exists,_it_fails", func(t *testing.T) { + s := setup(t) + req := &session.CreateRequest{ + AppName: testAppName, + UserID: "testUserID", + } + + got1, err := s.Create(t.Context(), req) + if err != nil { + t.Fatalf("First Create() failed: %v", err) + } + + req2 := &session.CreateRequest{ + AppName: req.AppName, + UserID: req.UserID, + SessionID: got1.Session.ID(), + } + + _, err = s.Create(t.Context(), req2) + if opts.SupportsUserProvidedSessionID { + if err == nil { + t.Errorf("Expected failure when creating duplicate session") + } + } else { + if err == nil { + t.Errorf("Expected failure (unsupported or duplicate)") + } + } + }) + }) + + t.Run("Get", func(t *testing.T) { + t.Run("ok", func(t *testing.T) { + s := setup(t) + req := &session.CreateRequest{ + AppName: testAppName, + UserID: "testUserID", + State: map[string]any{"k1": "v1"}, + } + created, err := s.Create(t.Context(), req) + if err != nil { + t.Fatalf("Setup: Create failed: %v", err) + } + + got, err := s.Get(t.Context(), &session.GetRequest{ + AppName: req.AppName, + UserID: req.UserID, + SessionID: created.Session.ID(), + }) + if err != nil { + t.Fatalf("Get() error = %v", err) + } + + snap := Snapshot(got.Session) + if snap.AppName != req.AppName { + t.Errorf("Get AppName = %v, want %v", snap.AppName, req.AppName) + } + if snap.UserID != req.UserID { + t.Errorf("Get UserID = %v, want %v", snap.UserID, req.UserID) + } + if snap.State["k1"] != "v1" { + t.Errorf("Get State[k1] = %v, want v1", snap.State["k1"]) + } + }) + + t.Run("error_when_not_found", func(t *testing.T) { + s := setup(t) + _, err := s.Get(t.Context(), &session.GetRequest{ + AppName: "nonExistent", + UserID: "user", + SessionID: "s1", + }) + if err == nil { + t.Errorf("Expected error for non-existent session") + } + }) + + t.Run("get_session_respects_user_id", func(t *testing.T) { + s := setup(t) + ctx := t.Context() + + c1, err := s.Create(ctx, &session.CreateRequest{AppName: testAppName, UserID: "user1"}) + if err != nil { + t.Fatalf("Create user1 failed: %v", err) + } + + err = s.AppendEvent(ctx, c1.Session, &session.Event{ + ID: "event1", + Author: "user", + InvocationID: "inv1", + }) + if err != nil { + t.Fatalf("AppendEvent failed: %v", err) + } + + _, err = s.Get(ctx, &session.GetRequest{ + AppName: testAppName, + UserID: "user2", + SessionID: c1.Session.ID(), + }) + if err == nil { + t.Errorf("Expected error or not found when getting session with wrong UserID") + } + }) + + t.Run("with_config_filters", func(t *testing.T) { + s := setup(t) + ctx := t.Context() + req := &session.CreateRequest{AppName: testAppName, UserID: "user1"} + if opts.SupportsUserProvidedSessionID { + req.SessionID = "s1" + } + created, err := s.Create(ctx, req) + if err != nil { + t.Fatalf("Setup: Create failed: %v", err) + } + + numTestEvents := 5 + var timestamps []time.Time + baseTime := time.Date(2025, 1, 1, 0, 0, 0, 0, time.UTC) + + for i := 1; i <= numTestEvents; i++ { + ts := baseTime.Add(time.Duration(i) * time.Second) + timestamps = append(timestamps, ts) + event := &session.Event{ + ID: strconv.Itoa(i), + Author: "user", + InvocationID: "inv1", + Timestamp: ts, + } + if err := s.AppendEvent(ctx, created.Session, event); err != nil { + t.Fatalf("AppendEvent failed: %v", err) + } + } + + gotNormal, err := s.Get(ctx, &session.GetRequest{AppName: testAppName, UserID: "user1", SessionID: created.Session.ID()}) + if err != nil { + t.Fatalf("Get error: %v", err) + } + if gotNormal.Session.Events().Len() != 5 { + t.Errorf("Get no filter events len = %d, want 5", gotNormal.Session.Events().Len()) + } + + gotLimit, err := s.Get(ctx, &session.GetRequest{AppName: testAppName, UserID: "user1", SessionID: created.Session.ID(), NumRecentEvents: 3}) + if err != nil { + t.Fatalf("Get error: %v", err) + } + if gotLimit.Session.Events().Len() != 3 { + t.Errorf("Get NumRecentEvents len = %d, want 3", gotLimit.Session.Events().Len()) + } + + gotAfter, err := s.Get(ctx, &session.GetRequest{AppName: testAppName, UserID: "user1", SessionID: created.Session.ID(), After: timestamps[1]}) + if err != nil { + t.Fatalf("Get error: %v", err) + } + if gotAfter.Session.Events().Len() != 4 { + t.Errorf("Get After filter events len = %d, want 4", gotAfter.Session.Events().Len()) + } + + gotCombined, err := s.Get(ctx, &session.GetRequest{AppName: testAppName, UserID: "user1", SessionID: created.Session.ID(), NumRecentEvents: 2, After: timestamps[0]}) + if err != nil { + t.Fatalf("Get error: %v", err) + } + if gotCombined.Session.Events().Len() != 2 { + t.Errorf("Get Combined filters len = %d, want 2", gotCombined.Session.Events().Len()) + } + }) + }) + + t.Run("List", func(t *testing.T) { + s := setup(t) + ctx := t.Context() + + _, err := s.Create(ctx, &session.CreateRequest{AppName: testAppName, UserID: "user1"}) + if err != nil { + t.Fatalf("Create user1 session1 failed: %v", err) + } + _, err = s.Create(ctx, &session.CreateRequest{AppName: testAppName, UserID: "user1"}) + if err != nil { + t.Fatalf("Create user1 session2 failed: %v", err) + } + + _, err = s.Create(ctx, &session.CreateRequest{AppName: testAppName, UserID: "user2"}) + if err != nil { + t.Fatalf("Create user2 failed: %v", err) + } + + got1, err := s.List(ctx, &session.ListRequest{AppName: testAppName, UserID: "user1"}) + if err != nil { + t.Fatalf("List user1 failed: %v", err) + } + if len(got1.Sessions) != 2 { + t.Errorf("List user1 len = %d, want 2", len(got1.Sessions)) + } + + got2, err := s.List(ctx, &session.ListRequest{AppName: testAppName, UserID: "user2"}) + if err != nil { + t.Fatalf("List user2 failed: %v", err) + } + if len(got2.Sessions) != 1 { + t.Errorf("List user2 len = %d, want 1", len(got2.Sessions)) + } + + got3, err := s.List(ctx, &session.ListRequest{AppName: testAppName, UserID: "nonExistent"}) + if err != nil { + t.Fatalf("List nonExistent failed: %v", err) + } + if len(got3.Sessions) != 0 { + t.Errorf("List nonExistent len = %d, want 0", len(got3.Sessions)) + } + + gotAll, err := s.List(ctx, &session.ListRequest{AppName: testAppName}) + if err != nil { + t.Fatalf("List all users failed: %v", err) + } + if len(gotAll.Sessions) != 3 { + t.Errorf("List all users len = %d, want 3", len(gotAll.Sessions)) + } + }) + + t.Run("Delete", func(t *testing.T) { + s := setup(t) + ctx := t.Context() + + created, err := s.Create(ctx, &session.CreateRequest{AppName: testAppName, UserID: "user1"}) + if err != nil { + t.Fatalf("Setup: Create failed: %v", err) + } + + err = s.Delete(ctx, &session.DeleteRequest{ + AppName: testAppName, + UserID: "user1", + SessionID: created.Session.ID(), + }) + if err != nil { + t.Fatalf("Delete() error = %v", err) + } + + _, err = s.Get(ctx, &session.GetRequest{ + AppName: testAppName, + UserID: "user1", + SessionID: created.Session.ID(), + }) + if err == nil { + t.Errorf("Expected error when getting deleted session") + } + + if opts.SupportsUserProvidedSessionID { + err = s.Delete(ctx, &session.DeleteRequest{ + AppName: testAppName, + UserID: "user1", + SessionID: "nonExistent", + }) + if err != nil { + if st, ok := status.FromError(err); ok { + // VertexAI fails on Delete when the sessionID is not found, we accept that + if st.Code() != codes.InvalidArgument { + t.Errorf("Delete() non-existent gRPC code = %v, want InvalidArgument", st.Code()) + } + } else { + t.Errorf("Delete() non-existent error = %v, want nil or gRPC InvalidArgument", err) + } + } + } + }) + + t.Run("AppendEvent", func(t *testing.T) { + t.Run("when_session_not_found_should_fail", func(t *testing.T) { + s := setup(t) + ctx := t.Context() + + m := &mockSession{appName: testAppName, userID: "user1", id: "nonExistent"} + event := &session.Event{ID: "event1", Author: "user", InvocationID: "inv1"} + + err := s.AppendEvent(ctx, m, event) + if err == nil { + t.Errorf("AppendEvent() expected error for non-existent session, got nil") + } + }) + + t.Run("ok", func(t *testing.T) { + s := setup(t) + ctx := t.Context() + + created, err := s.Create(ctx, &session.CreateRequest{AppName: testAppName, UserID: "user1"}) + if err != nil { + t.Fatalf("Setup: Create failed: %v", err) + } + + event := &session.Event{ + ID: "new_event1", + Author: "user", + InvocationID: "inv1", + } + err = s.AppendEvent(ctx, created.Session, event) + if err != nil { + t.Fatalf("AppendEvent() error = %v", err) + } + + got, err := s.Get(ctx, &session.GetRequest{ + AppName: testAppName, + UserID: "user1", + SessionID: created.Session.ID(), + }) + if err != nil { + t.Fatalf("Get() error = %v", err) + } + + snap := Snapshot(got.Session) + if len(snap.Events) != 1 { + t.Errorf("Expected 1 event, got %d", len(snap.Events)) + } else { + if !opts.ProvidesServerAssignedEventID && snap.Events[0].ID != event.ID { + t.Errorf("Event ID mismatch: got %v, want %v", snap.Events[0].ID, event.ID) + } + } + }) + + t.Run("partial_events_are_not_persisted", func(t *testing.T) { + s := setup(t) + ctx := t.Context() + + created, err := s.Create(ctx, &session.CreateRequest{AppName: testAppName, UserID: "user1"}) + if err != nil { + t.Fatalf("Setup: Create failed: %v", err) + } + + event := &session.Event{ + ID: "partial_event", + Author: "user", + InvocationID: "inv1", + LLMResponse: model.LLMResponse{ + Partial: true, + }, + } + err = s.AppendEvent(ctx, created.Session, event) + if err != nil { + t.Fatalf("AppendEvent() error = %v", err) + } + + got, err := s.Get(ctx, &session.GetRequest{ + AppName: testAppName, + UserID: "user1", + SessionID: created.Session.ID(), + }) + if err != nil { + t.Fatalf("Get() error = %v", err) + } + + snap := Snapshot(got.Session) + if len(snap.Events) != 0 { + t.Errorf("Expected 0 events (Partial=true should be skipped), got %d", len(snap.Events)) + } + }) + + t.Run("with_bytes_content", func(t *testing.T) { + s := setup(t) + ctx := t.Context() + + created, err := s.Create(ctx, &session.CreateRequest{AppName: testAppName, UserID: "user1"}) + if err != nil { + t.Fatalf("Setup: Create failed: %v", err) + } + + event := &session.Event{ + ID: "event_with_bytes", + Author: "user", + InvocationID: "inv1", + LLMResponse: model.LLMResponse{ + Content: genai.NewContentFromBytes([]byte("test_image_data"), "image/png", "user"), + }, + } + err = s.AppendEvent(ctx, created.Session, event) + if err != nil { + t.Fatalf("AppendEvent() error = %v", err) + } + + got, err := s.Get(ctx, &session.GetRequest{ + AppName: testAppName, + UserID: "user1", + SessionID: created.Session.ID(), + }) + if err != nil { + t.Fatalf("Get() error = %v", err) + } + + snap := Snapshot(got.Session) + if len(snap.Events) != 1 { + t.Fatalf("Expected 1 event, got %d", len(snap.Events)) + } + gotEvent := snap.Events[0] + + if diff := cmp.Diff(event.LLMResponse.Content, gotEvent.LLMResponse.Content); diff != "" { + t.Errorf("Content mismatch (-want +got):\n%s", diff) + } + }) + + t.Run("with_existing_events", func(t *testing.T) { + s := setup(t) + ctx := t.Context() + + created, err := s.Create(ctx, &session.CreateRequest{AppName: testAppName, UserID: "user1"}) + if err != nil { + t.Fatalf("Setup: Create failed: %v", err) + } + + baseTime := time.Date(2025, 1, 1, 0, 0, 0, 0, time.UTC) + event1 := &session.Event{ID: "event1", Author: "user", InvocationID: "inv1", Timestamp: baseTime.Add(1 * time.Second)} + if err := s.AppendEvent(ctx, created.Session, event1); err != nil { + t.Fatalf("AppendEvent(1) failed: %v", err) + } + + event2 := &session.Event{ID: "event2", Author: "user", InvocationID: "inv2", Timestamp: baseTime.Add(2 * time.Second)} + if err := s.AppendEvent(ctx, created.Session, event2); err != nil { + t.Fatalf("AppendEvent(2) failed: %v", err) + } + + got, err := s.Get(ctx, &session.GetRequest{ + AppName: testAppName, + UserID: "user1", + SessionID: created.Session.ID(), + }) + if err != nil { + t.Fatalf("Get() failed: %v", err) + } + + snap := Snapshot(got.Session) + if len(snap.Events) != 2 { + t.Fatalf("Expected 2 events, got %d", len(snap.Events)) + } + expectedOrder := []string{"inv1", "inv2"} + for i := range snap.Events { + if snap.Events[i].InvocationID != expectedOrder[i] { + t.Errorf("Expected event order for %v, got events with InvocationIDs", expectedOrder) + break + } + } + }) + + t.Run("with_all_fields", func(t *testing.T) { + s := setup(t) + ctx := t.Context() + + created, err := s.Create(ctx, &session.CreateRequest{AppName: testAppName, UserID: "user1"}) + if err != nil { + t.Fatalf("Setup: Create failed: %v", err) + } + + event := &session.Event{ + ID: "event_complete", + Author: "user", + InvocationID: "inv1", + LongRunningToolIDs: []string{"tool123"}, + Actions: session.EventActions{StateDelta: map[string]any{"user:k2": "v2"}}, + LLMResponse: model.LLMResponse{ + Content: genai.NewContentFromText("test_text", "user"), + TurnComplete: true, + Partial: false, + ErrorCode: "error_code", + ErrorMessage: "error_message", + Interrupted: true, + GroundingMetadata: &genai.GroundingMetadata{ + WebSearchQueries: []string{"query1"}, + }, + UsageMetadata: &genai.GenerateContentResponseUsageMetadata{ + PromptTokenCount: 1, + CandidatesTokenCount: 1, + TotalTokenCount: 2, + }, + CitationMetadata: &genai.CitationMetadata{ + Citations: []*genai.Citation{{Title: "test", URI: "google.com"}}, + }, + CustomMetadata: map[string]any{ + "custom_key": "custom_value", + }, + }, + } + + err = s.AppendEvent(ctx, created.Session, event) + if err != nil { + t.Fatalf("AppendEvent() error = %v", err) + } + + got, err := s.Get(ctx, &session.GetRequest{ + AppName: testAppName, + UserID: "user1", + SessionID: created.Session.ID(), + }) + if err != nil { + t.Fatalf("Get() error = %v", err) + } + + snap := Snapshot(got.Session) + if len(snap.Events) != 1 { + t.Fatalf("Expected 1 event, got %d", len(snap.Events)) + } + gotEvent := snap.Events[0] + + cmpOpts := []cmp.Option{ + cmp.AllowUnexported(session.Event{}), + } + if opts.ProvidesServerAssignedEventID { + cmpOpts = append(cmpOpts, cmpopts.IgnoreFields(session.Event{}, "ID")) + cmpOpts = append(cmpOpts, cmpopts.IgnoreFields(model.LLMResponse{}, "CitationMetadata", "UsageMetadata")) + } + + if diff := cmp.Diff(event, gotEvent, cmpOpts...); diff != "" { + t.Errorf("Event mismatch (-want +got):\n%s", diff) + } + }) + }) + + t.Run("StateManagement", func(t *testing.T) { + ctx := t.Context() + appName := testAppName + + t.Run("app_state_is_shared", func(t *testing.T) { + s := setup(t) + s1, _ := s.Create(ctx, &session.CreateRequest{AppName: appName, UserID: "u1", State: map[string]any{"app:k1": "v1"}}) + if err := s.AppendEvent(ctx, s1.Session, &session.Event{ + ID: "event1", + Author: "user", + InvocationID: "inv1", + Actions: session.EventActions{StateDelta: map[string]any{"app:k2": "v2"}}, + }); err != nil { + t.Fatalf("AppendEvent failed: %v", err) + } + + s2, err := s.Create(ctx, &session.CreateRequest{AppName: appName, UserID: "u2"}) + if err != nil { + t.Fatalf("Create user2 failed: %v", err) + } + + snap := Snapshot(s2.Session) + if snap.State["app:k1"] != "v1" || snap.State["app:k2"] != "v2" { + t.Errorf("App state not shared, got: %v", snap.State) + } + }) + + t.Run("user_state_is_user_specific", func(t *testing.T) { + s := setup(t) + s1, _ := s.Create(ctx, &session.CreateRequest{AppName: appName, UserID: "u1", State: map[string]any{"user:k1": "v1"}}) + if err := s.AppendEvent(ctx, s1.Session, &session.Event{ + ID: "event1", + Author: "user", + InvocationID: "inv1", + Actions: session.EventActions{StateDelta: map[string]any{"user:k2": "v2"}}, + }); err != nil { + t.Fatalf("AppendEvent failed: %v", err) + } + + s1b, _ := s.Create(ctx, &session.CreateRequest{AppName: appName, UserID: "u1"}) + snap1b := Snapshot(s1b.Session) + if snap1b.State["user:k1"] != "v1" || snap1b.State["user:k2"] != "v2" { + t.Errorf("User state missing for same user, got: %v", snap1b.State) + } + + s2, _ := s.Create(ctx, &session.CreateRequest{AppName: appName, UserID: "u2"}) + snap2 := Snapshot(s2.Session) + if _, exists := snap2.State["user:k1"]; exists { + t.Errorf("User state leaked to user2, got: %v", snap2.State) + } + }) + + t.Run("session_state_is_not_shared", func(t *testing.T) { + s := setup(t) + s1, _ := s.Create(ctx, &session.CreateRequest{AppName: appName, UserID: "u1", State: map[string]any{"sk1": "v1"}}) + if err := s.AppendEvent(ctx, s1.Session, &session.Event{ + ID: "event1", + Author: "user", + InvocationID: "inv1", + Actions: session.EventActions{StateDelta: map[string]any{"sk2": "v2"}}, + }); err != nil { + t.Fatalf("AppendEvent failed: %v", err) + } + + s1b, _ := s.Create(ctx, &session.CreateRequest{AppName: appName, UserID: "u1"}) + snapS1b := Snapshot(s1b.Session) + if _, exists := snapS1b.State["sk1"]; exists { + t.Errorf("Session state leaked between sessions, got: %v", snapS1b.State) + } + }) + + t.Run("temp_state_is_not_persisted", func(t *testing.T) { + s := setup(t) + s1, _ := s.Create(ctx, &session.CreateRequest{AppName: appName, UserID: "u1"}) + _ = s.AppendEvent(ctx, s1.Session, &session.Event{ + ID: "event1", + Author: "user", + InvocationID: "inv1", + Actions: session.EventActions{StateDelta: map[string]any{"temp:k1": "v1", "sk": "v2"}}, + }) + + got, _ := s.Get(ctx, &session.GetRequest{AppName: appName, UserID: "u1", SessionID: s1.Session.ID()}) + snap := Snapshot(got.Session) + if _, exists := snap.State["temp:k1"]; exists { + t.Errorf("Temp state leaked to persist step, got: %v", snap.State) + } + if snap.State["sk"] != "v2" { + t.Errorf("Standard state update missing: got %v", snap.State) + } + }) + }) +} + +type mockSession struct { + appName string + userID string + id string +} + +func (m *mockSession) ID() string { return m.id } +func (m *mockSession) AppName() string { return m.appName } +func (m *mockSession) UserID() string { return m.userID } +func (m *mockSession) State() session.State { return nil } +func (m *mockSession) Events() session.Events { return nil } +func (m *mockSession) LastUpdateTime() time.Time { return time.Now() } diff --git a/session/vertexai/service_test.go b/session/vertexai/service_test.go index 0f5286818..2ce87faf4 100644 --- a/session/vertexai/service_test.go +++ b/session/vertexai/service_test.go @@ -28,24 +28,30 @@ import ( "google.golang.org/grpc" "google.golang.org/adk/session" - "google.golang.org/adk/session/session_test" + "google.golang.org/adk/session/sessiontestsuite" ) +// if you want to test it for yourself, you can regenerate all by running +// `UPDATE_REPLAYS=true go test --run ./... -v` +// You need a GCP project with enabled Agent Platform API +// and some deployed Agent Engine instances +// To deploy an instance of an application to Agent Engine you can run +// `go run ./cmd/adkgo/adkgo.go deploy agentengine -e ./examples/agentengine/main.go -p YOUR-PROJECT -r us-central1 -d . -s "Test01" ` const ( - ProjectID = "adk-go-test" + ProjectID = "adk-go-e2e" Location = "us-central1" - EngineId = "5576569044451983360" - EngineId2 = "8602987994044956672" + EngineID = "1491331942182813696" + EngineID2 = "6857370898194759680" UserID = "test-user" ) func Test_vertexaiService(t *testing.T) { - opts := session_test.SuiteOptions{ - SupportsUserProvidedSessionID: false, + opts := sessiontestsuite.SuiteOptions{ + SupportsUserProvidedSessionID: true, ProvidesServerAssignedEventID: true, - AppName: EngineId, + AppName: EngineID, } // VertexAI forbids custom IDs - session_test.RunServiceTests(t, opts, func(t *testing.T) session.Service { + sessiontestsuite.RunServiceTests(t, opts, func(t *testing.T) session.Service { name := strings.ReplaceAll(t.Name(), "/", "_") s, _ := emptyService(t, name, false) return s @@ -62,21 +68,21 @@ func Test_vertexaiService_AppendEvent_StructuralValidation(t *testing.T) { }{ { name: "missing_session_id", - session: &localSession{appName: EngineId, userID: UserID}, + session: &localSession{appName: EngineID, userID: UserID}, event: &session.Event{}, wantErr: true, offline: true, }, { name: "nil_event", - session: &localSession{appName: EngineId2, userID: "user2", sessionID: "session2"}, + session: &localSession{appName: EngineID2, userID: "user2", sessionID: "session2"}, event: nil, wantErr: true, offline: true, }, { name: "missing_author", - session: &localSession{appName: EngineId2, userID: "user2", sessionID: "session2"}, + session: &localSession{appName: EngineID2, userID: "user2", sessionID: "session2"}, event: &session.Event{ Timestamp: time.Now(), InvocationID: uuid.NewString(), @@ -85,7 +91,7 @@ func Test_vertexaiService_AppendEvent_StructuralValidation(t *testing.T) { }, { name: "missing_invocation_id", - session: &localSession{appName: EngineId2, userID: "user2", sessionID: "session2"}, + session: &localSession{appName: EngineID2, userID: "user2", sessionID: "session2"}, event: &session.Event{ Timestamp: time.Now(), Author: UserID, @@ -147,8 +153,8 @@ func emptyService(t *testing.T, name string, offline bool) (session.Service, map } func deleteAll(t *testing.T, v session.Service) { - deleteAllFromApp(t, v, EngineId) - deleteAllFromApp(t, v, EngineId2) + deleteAllFromApp(t, v, EngineID) + deleteAllFromApp(t, v, EngineID2) } func deleteAllFromApp(t *testing.T, v session.Service, app string) { @@ -158,6 +164,7 @@ func deleteAllFromApp(t *testing.T, v session.Service, app string) { }) if err != nil { t.Errorf("error listing session for delete all: %s", err) + return } for _, s := range sessionsResp.Sessions { diff --git a/session/vertexai/testdata/Test_vertexaiService_AppendEvent_ok.replay b/session/vertexai/testdata/Test_vertexaiService_AppendEvent_ok.replay index 743e91c35..eaf302fd6 100644 Binary files a/session/vertexai/testdata/Test_vertexaiService_AppendEvent_ok.replay and b/session/vertexai/testdata/Test_vertexaiService_AppendEvent_ok.replay differ diff --git a/session/vertexai/testdata/Test_vertexaiService_AppendEvent_partial_events_are_not_persisted.replay b/session/vertexai/testdata/Test_vertexaiService_AppendEvent_partial_events_are_not_persisted.replay index 6ffdf6b88..54f6fbf21 100644 Binary files a/session/vertexai/testdata/Test_vertexaiService_AppendEvent_partial_events_are_not_persisted.replay and b/session/vertexai/testdata/Test_vertexaiService_AppendEvent_partial_events_are_not_persisted.replay differ diff --git a/session/vertexai/testdata/Test_vertexaiService_AppendEvent_when_session_not_found_should_fail.replay b/session/vertexai/testdata/Test_vertexaiService_AppendEvent_when_session_not_found_should_fail.replay index da1ac3c0e..9edebdf48 100644 Binary files a/session/vertexai/testdata/Test_vertexaiService_AppendEvent_when_session_not_found_should_fail.replay and b/session/vertexai/testdata/Test_vertexaiService_AppendEvent_when_session_not_found_should_fail.replay differ diff --git a/session/vertexai/testdata/Test_vertexaiService_AppendEvent_with_all_fields.replay b/session/vertexai/testdata/Test_vertexaiService_AppendEvent_with_all_fields.replay index 42dab8697..f36bca646 100644 Binary files a/session/vertexai/testdata/Test_vertexaiService_AppendEvent_with_all_fields.replay and b/session/vertexai/testdata/Test_vertexaiService_AppendEvent_with_all_fields.replay differ diff --git a/session/vertexai/testdata/Test_vertexaiService_AppendEvent_with_bytes_content.replay b/session/vertexai/testdata/Test_vertexaiService_AppendEvent_with_bytes_content.replay index a00366e2f..718b28761 100644 Binary files a/session/vertexai/testdata/Test_vertexaiService_AppendEvent_with_bytes_content.replay and b/session/vertexai/testdata/Test_vertexaiService_AppendEvent_with_bytes_content.replay differ diff --git a/session/vertexai/testdata/Test_vertexaiService_AppendEvent_with_existing_events.replay b/session/vertexai/testdata/Test_vertexaiService_AppendEvent_with_existing_events.replay index b6f1b97e2..124b0a307 100644 Binary files a/session/vertexai/testdata/Test_vertexaiService_AppendEvent_with_existing_events.replay and b/session/vertexai/testdata/Test_vertexaiService_AppendEvent_with_existing_events.replay differ diff --git a/session/vertexai/testdata/Test_vertexaiService_Create_full_key.replay b/session/vertexai/testdata/Test_vertexaiService_Create_full_key.replay new file mode 100644 index 000000000..e716abd0a Binary files /dev/null and b/session/vertexai/testdata/Test_vertexaiService_Create_full_key.replay differ diff --git a/session/vertexai/testdata/Test_vertexaiService_Create_generated_session_id.replay b/session/vertexai/testdata/Test_vertexaiService_Create_generated_session_id.replay index 707357cc3..f85ca3c50 100644 Binary files a/session/vertexai/testdata/Test_vertexaiService_Create_generated_session_id.replay and b/session/vertexai/testdata/Test_vertexaiService_Create_generated_session_id.replay differ diff --git a/session/vertexai/testdata/Test_vertexaiService_Create_when_already_exists__it_fails.replay b/session/vertexai/testdata/Test_vertexaiService_Create_when_already_exists__it_fails.replay index 9628115f3..937f8d8c8 100644 Binary files a/session/vertexai/testdata/Test_vertexaiService_Create_when_already_exists__it_fails.replay and b/session/vertexai/testdata/Test_vertexaiService_Create_when_already_exists__it_fails.replay differ diff --git a/session/vertexai/testdata/Test_vertexaiService_Delete.replay b/session/vertexai/testdata/Test_vertexaiService_Delete.replay index 81f37ac61..00f0545fd 100644 Binary files a/session/vertexai/testdata/Test_vertexaiService_Delete.replay and b/session/vertexai/testdata/Test_vertexaiService_Delete.replay differ diff --git a/session/vertexai/testdata/Test_vertexaiService_Get_error_when_not_found.replay b/session/vertexai/testdata/Test_vertexaiService_Get_error_when_not_found.replay index 71f8ef74f..14702fed7 100644 Binary files a/session/vertexai/testdata/Test_vertexaiService_Get_error_when_not_found.replay and b/session/vertexai/testdata/Test_vertexaiService_Get_error_when_not_found.replay differ diff --git a/session/vertexai/testdata/Test_vertexaiService_Get_get_session_respects_user_id.replay b/session/vertexai/testdata/Test_vertexaiService_Get_get_session_respects_user_id.replay index f4615a8a1..eec25c86e 100644 Binary files a/session/vertexai/testdata/Test_vertexaiService_Get_get_session_respects_user_id.replay and b/session/vertexai/testdata/Test_vertexaiService_Get_get_session_respects_user_id.replay differ diff --git a/session/vertexai/testdata/Test_vertexaiService_Get_ok.replay b/session/vertexai/testdata/Test_vertexaiService_Get_ok.replay index 386b3a6fa..b728c9f0a 100644 Binary files a/session/vertexai/testdata/Test_vertexaiService_Get_ok.replay and b/session/vertexai/testdata/Test_vertexaiService_Get_ok.replay differ diff --git a/session/vertexai/testdata/Test_vertexaiService_Get_with_config_filters.replay b/session/vertexai/testdata/Test_vertexaiService_Get_with_config_filters.replay index 29fd3132a..2e921064e 100644 Binary files a/session/vertexai/testdata/Test_vertexaiService_Get_with_config_filters.replay and b/session/vertexai/testdata/Test_vertexaiService_Get_with_config_filters.replay differ diff --git a/session/vertexai/testdata/Test_vertexaiService_List.replay b/session/vertexai/testdata/Test_vertexaiService_List.replay index a41624b31..df1269e44 100644 Binary files a/session/vertexai/testdata/Test_vertexaiService_List.replay and b/session/vertexai/testdata/Test_vertexaiService_List.replay differ diff --git a/session/vertexai/testdata/Test_vertexaiService_StateManagement_app_state_is_shared.replay b/session/vertexai/testdata/Test_vertexaiService_StateManagement_app_state_is_shared.replay index 721a5c7ba..fd114798b 100644 Binary files a/session/vertexai/testdata/Test_vertexaiService_StateManagement_app_state_is_shared.replay and b/session/vertexai/testdata/Test_vertexaiService_StateManagement_app_state_is_shared.replay differ diff --git a/session/vertexai/testdata/Test_vertexaiService_StateManagement_session_state_is_not_shared.replay b/session/vertexai/testdata/Test_vertexaiService_StateManagement_session_state_is_not_shared.replay index e0f8c5499..d263387e1 100644 Binary files a/session/vertexai/testdata/Test_vertexaiService_StateManagement_session_state_is_not_shared.replay and b/session/vertexai/testdata/Test_vertexaiService_StateManagement_session_state_is_not_shared.replay differ diff --git a/session/vertexai/testdata/Test_vertexaiService_StateManagement_temp_state_is_not_persisted.replay b/session/vertexai/testdata/Test_vertexaiService_StateManagement_temp_state_is_not_persisted.replay index 9046e20a8..2a0747251 100644 Binary files a/session/vertexai/testdata/Test_vertexaiService_StateManagement_temp_state_is_not_persisted.replay and b/session/vertexai/testdata/Test_vertexaiService_StateManagement_temp_state_is_not_persisted.replay differ diff --git a/session/vertexai/testdata/Test_vertexaiService_StateManagement_user_state_is_user_specific.replay b/session/vertexai/testdata/Test_vertexaiService_StateManagement_user_state_is_user_specific.replay index c37fdba64..809035f5a 100644 Binary files a/session/vertexai/testdata/Test_vertexaiService_StateManagement_user_state_is_user_specific.replay and b/session/vertexai/testdata/Test_vertexaiService_StateManagement_user_state_is_user_specific.replay differ diff --git a/session/vertexai/testdata/missing_author.replay b/session/vertexai/testdata/missing_author.replay index eed191ab6..fd66f499e 100644 Binary files a/session/vertexai/testdata/missing_author.replay and b/session/vertexai/testdata/missing_author.replay differ diff --git a/session/vertexai/testdata/missing_invocation_id.replay b/session/vertexai/testdata/missing_invocation_id.replay index 83c71691c..5437eb2f5 100644 Binary files a/session/vertexai/testdata/missing_invocation_id.replay and b/session/vertexai/testdata/missing_invocation_id.replay differ diff --git a/session/vertexai/vertexai.go b/session/vertexai/vertexai.go index b7b0bb523..5af33944d 100644 --- a/session/vertexai/vertexai.go +++ b/session/vertexai/vertexai.go @@ -56,9 +56,6 @@ func (s *vertexAiService) Create(ctx context.Context, req *session.CreateRequest if req.AppName == "" || req.UserID == "" { return nil, fmt.Errorf("app_name and user_id are required, got app_name: %q, user_id: %q", req.AppName, req.UserID) } - if req.SessionID != "" { - return nil, fmt.Errorf("user-provided Session id is not supported for VertexAISessionService: %q", req.SessionID) - } sess, err := s.client.createSession(ctx, req) if err != nil { return nil, fmt.Errorf("failed to create session: %w", err) diff --git a/session/vertexai/vertexai_client.go b/session/vertexai/vertexai_client.go index 04f1b5000..b31724da3 100644 --- a/session/vertexai/vertexai_client.go +++ b/session/vertexai/vertexai_client.go @@ -16,6 +16,7 @@ package vertexai import ( "context" + "encoding/json" "fmt" "regexp" "strconv" @@ -32,20 +33,14 @@ import ( "google.golang.org/adk/model" "google.golang.org/adk/session" + vertexaiutil "google.golang.org/adk/util/vertexai" aiplatform "cloud.google.com/go/aiplatform/apiv1beta1" aiplatformpb "cloud.google.com/go/aiplatform/apiv1beta1/aiplatformpb" ) -const ( - engineResourceTemplate = "projects/%s/locations/%s/reasoningEngines/%s" - sessionResourceTemplate = engineResourceTemplate + "/sessions/%s" -) - type vertexAiClient struct { - location string - projectID string - reasoningEngine string + agentEngineData *vertexaiutil.AgentEngineData rpcClient *aiplatform.SessionClient } @@ -54,7 +49,14 @@ func newVertexAiClient(ctx context.Context, location, projectID, reasoningEngine if err != nil { return nil, fmt.Errorf("could not establish connection to the aiplatform server: %w", err) } - return &vertexAiClient{location, projectID, reasoningEngine, rpcClient}, nil + return &vertexAiClient{ + agentEngineData: &vertexaiutil.AgentEngineData{ + Location: location, + ProjectID: projectID, + ReasoningEngine: reasoningEngine, + }, + rpcClient: rpcClient, + }, nil } // Ensure you close it when your application shuts down @@ -68,7 +70,7 @@ func (c *vertexAiClient) createSession(ctx context.Context, req *session.CreateR } // Convert and set the initial state if provided if len(req.State) > 0 { - stateStruct, err := structpb.NewStruct(req.State) + stateStruct, err := toStructPB(req.State) if err != nil { return nil, fmt.Errorf("failed to convert state to structpb: %w", err) } @@ -79,9 +81,15 @@ func (c *vertexAiClient) createSession(ctx context.Context, req *session.CreateR if err != nil { return nil, err } + aeData := &vertexaiutil.AgentEngineData{ + Location: c.agentEngineData.Location, + ProjectID: c.agentEngineData.ProjectID, + ReasoningEngine: reasoningEngine, + } rpcReq := &aiplatformpb.CreateSessionRequest{ - Parent: fmt.Sprintf(engineResourceTemplate, c.projectID, c.location, reasoningEngine), - Session: pbSession, + Parent: vertexaiutil.AgentEngineResource(aeData), + Session: pbSession, + SessionId: req.SessionID, } lro, err := c.rpcClient.CreateSession(ctx, rpcReq) if err != nil { @@ -168,8 +176,15 @@ func (c *vertexAiClient) listSessions(ctx context.Context, req *session.ListRequ if err != nil { return nil, err } + + aeData := vertexaiutil.AgentEngineData{ + Location: c.agentEngineData.Location, + ProjectID: c.agentEngineData.ProjectID, + ReasoningEngine: reasoningEngine, + } + rpcReq := &aiplatformpb.ListSessionsRequest{ - Parent: fmt.Sprintf(engineResourceTemplate, c.projectID, c.location, reasoningEngine), + Parent: vertexaiutil.AgentEngineResource(&aeData), } if req.UserID != "" { rpcReq.Filter = fmt.Sprintf("userId=\"%s\"", req.UserID) @@ -241,7 +256,7 @@ func (c *vertexAiClient) appendEvent(ctx context.Context, appName, sessionID str var eventState *aiplatformpb.EventActions // Convert and set the initial state if provided if len(event.Actions.StateDelta) > 0 { - sessionState, err := structpb.NewStruct(event.Actions.StateDelta) + sessionState, err := toStructPB(event.Actions.StateDelta) if err != nil { return fmt.Errorf("failed to convert state to structpb: %w", err) } @@ -392,7 +407,12 @@ func sessionIDByOperationName(on string) (string, error) { } func sessionNameByID(id string, c *vertexAiClient, reasoningEngine string) string { - return fmt.Sprintf(sessionResourceTemplate, c.projectID, c.location, reasoningEngine, id) + aeData := &vertexaiutil.AgentEngineData{ + Location: c.agentEngineData.Location, + ProjectID: c.agentEngineData.ProjectID, + ReasoningEngine: reasoningEngine, + } + return vertexaiutil.SessionResource(aeData, id) } // (?:...) tells Go "match this, but don't save it in the results array". @@ -400,8 +420,8 @@ func sessionNameByID(id string, c *vertexAiClient, reasoningEngine string) strin var reasoningEnginePattern = regexp.MustCompile(`^projects/(?:[a-zA-Z0-9-_]+)/locations/(?:[a-zA-Z0-9-_]+)/reasoningEngines/(\d+)$`) func (c *vertexAiClient) getReasoningEngineID(appName string) (string, error) { - if c.reasoningEngine != "" { - return c.reasoningEngine, nil + if c.agentEngineData.ReasoningEngine != "" { + return c.agentEngineData.ReasoningEngine, nil } // Check if appName consists only of digits @@ -443,12 +463,14 @@ func aiplatformToGenaiContent(rpcResp *aiplatformpb.SessionEvent) *genai.Content case *aiplatformpb.Part_FunctionCall: argsMap := v.FunctionCall.Args.AsMap() // Converts *structpb.Struct -> map[string]any part.FunctionCall = &genai.FunctionCall{ + ID: v.FunctionCall.Id, Name: v.FunctionCall.Name, Args: argsMap, } case *aiplatformpb.Part_FunctionResponse: responseMap := v.FunctionResponse.Response.AsMap() // Converts *structpb.Struct -> map[string]any part.FunctionResponse = &genai.FunctionResponse{ + ID: v.FunctionResponse.Id, Name: v.FunctionResponse.Name, Response: responseMap, } @@ -484,7 +506,7 @@ func createAiplatformpbContent(event *session.Event) (*aiplatformpb.Content, err } } if part.FunctionCall != nil { - args, err := structpb.NewStruct(part.FunctionCall.Args) + args, err := toStructPB(part.FunctionCall.Args) if err != nil { return nil, fmt.Errorf("failed to convert function call to structpb: %w", err) } @@ -497,7 +519,7 @@ func createAiplatformpbContent(event *session.Event) (*aiplatformpb.Content, err } } if part.FunctionResponse != nil { - response, err := structpb.NewStruct(part.FunctionResponse.Response) + response, err := toStructPB(part.FunctionResponse.Response) if err != nil { return nil, fmt.Errorf("failed to convert function response to structpb: %w", err) } @@ -531,7 +553,7 @@ func createAiplatformpbMetadata(event *session.Event) (*aiplatformpb.EventMetada Branch: event.Branch, } if event.CustomMetadata != nil { - customMetadata, err := structpb.NewStruct(event.CustomMetadata) + customMetadata, err := toStructPB(event.CustomMetadata) if err != nil { return nil, fmt.Errorf("failed to convert event customMetadata to structpb: %w", err) } @@ -763,6 +785,22 @@ func createGroundingMetadata(metadata *aiplatformpb.GroundingMetadata) *genai.Gr return out } +// toStructPB converts an arbitrary Go value into a protobuf Struct. +// It uses JSON marshaling as an intermediary step to safely serialize +// the input data before constructing the *structpb.Struct. +// Returns an error if any part of the JSON round-trip or conversion fails. +func toStructPB(value any) (*structpb.Struct, error) { + data, err := json.Marshal(value) + if err != nil { + return nil, fmt.Errorf("failed to marshal value: %w", err) + } + res := &structpb.Struct{} + if err := res.UnmarshalJSON(data); err != nil { + return nil, fmt.Errorf("failed to unmarshal JSON data to structpb: %w", err) + } + return res, nil +} + // derefString is a helper to safely dereference string pointers // Returns empty string if pointer is nil func derefString(s *string) string { diff --git a/session/vertexai/vertexai_test.go b/session/vertexai/vertexai_test.go index 2e4b9251a..9b37a5fa8 100644 --- a/session/vertexai/vertexai_test.go +++ b/session/vertexai/vertexai_test.go @@ -16,8 +16,147 @@ package vertexai import ( "testing" + + "google.golang.org/adk/model" + "google.golang.org/adk/session" + "google.golang.org/adk/util/vertexai" + + aiplatformpb "cloud.google.com/go/aiplatform/apiv1beta1/aiplatformpb" + "google.golang.org/genai" + "google.golang.org/protobuf/types/known/structpb" ) +func TestAiplatformToGenaiContent_FunctionCallMapping(t *testing.T) { + makeArgs := func(m map[string]any) *structpb.Struct { + s, err := structpb.NewStruct(m) + if err != nil { + t.Fatalf("failed to create struct: %v", err) + } + return s + } + + tests := []struct { + name string + input *aiplatformpb.SessionEvent + wantID string + wantName string + wantArgKey string + wantArgVal string + isResponse bool + wantRespKey string + wantRespVal string + }{ + { + name: "FunctionCall preserves ID, Name, and Args", + input: &aiplatformpb.SessionEvent{ + Content: &aiplatformpb.Content{ + Role: "model", + Parts: []*aiplatformpb.Part{ + { + Data: &aiplatformpb.Part_FunctionCall{ + FunctionCall: &aiplatformpb.FunctionCall{ + Id: "call-id-abc", + Name: "my_tool", + Args: makeArgs(map[string]any{"param": "value"}), + }, + }, + }, + }, + }, + }, + wantID: "call-id-abc", + wantName: "my_tool", + wantArgKey: "param", + wantArgVal: "value", + }, + { + name: "FunctionCall with empty ID is preserved as empty", + input: &aiplatformpb.SessionEvent{ + Content: &aiplatformpb.Content{ + Role: "model", + Parts: []*aiplatformpb.Part{ + { + Data: &aiplatformpb.Part_FunctionCall{ + FunctionCall: &aiplatformpb.FunctionCall{ + Id: "", + Name: "tool_no_id", + Args: makeArgs(map[string]any{"x": "y"}), + }, + }, + }, + }, + }, + }, + wantID: "", + wantName: "tool_no_id", + wantArgKey: "x", + wantArgVal: "y", + }, + { + name: "FunctionResponse preserves ID, Name, and Response", + isResponse: true, + input: &aiplatformpb.SessionEvent{ + Content: &aiplatformpb.Content{ + Role: "user", + Parts: []*aiplatformpb.Part{ + { + Data: &aiplatformpb.Part_FunctionResponse{ + FunctionResponse: &aiplatformpb.FunctionResponse{ + Id: "call-id-abc", + Name: "my_tool", + Response: makeArgs(map[string]any{"result": "ok"}), + }, + }, + }, + }, + }, + }, + wantID: "call-id-abc", + wantName: "my_tool", + wantRespKey: "result", + wantRespVal: "ok", + }, + } + + for _, tt := range tests { + t.Run(tt.name, func(t *testing.T) { + got := aiplatformToGenaiContent(tt.input) + if got == nil || len(got.Parts) == 0 { + t.Fatal("expected at least one part, got nil or empty") + } + if tt.isResponse { + fr := got.Parts[0].FunctionResponse + if fr == nil { + t.Fatal("expected FunctionResponse part, got nil") + } + if fr.ID != tt.wantID { + t.Errorf("FunctionResponse.ID = %q, want %q", fr.ID, tt.wantID) + } + if fr.Name != tt.wantName { + t.Errorf("FunctionResponse.Name = %q, want %q", fr.Name, tt.wantName) + } + if got, ok := fr.Response[tt.wantRespKey]; !ok || got != tt.wantRespVal { + t.Errorf("FunctionResponse.Response[%q] = %v, want %q", tt.wantRespKey, got, tt.wantRespVal) + } + } else { + fc := got.Parts[0].FunctionCall + if fc == nil { + t.Fatal("expected FunctionCall part, got nil") + } + if fc.ID != tt.wantID { + t.Errorf("FunctionCall.ID = %q, want %q", fc.ID, tt.wantID) + } + if fc.Name != tt.wantName { + t.Errorf("FunctionCall.Name = %q, want %q", fc.Name, tt.wantName) + } + if got, ok := fc.Args[tt.wantArgKey]; !ok || got != tt.wantArgVal { + t.Errorf("FunctionCall.Args[%q] = %v, want %q", tt.wantArgKey, got, tt.wantArgVal) + } + } + }) + } +} + func TestGetReasoningEngineID(t *testing.T) { tests := []struct { name string @@ -81,7 +220,9 @@ func TestGetReasoningEngineID(t *testing.T) { t.Run(tt.name, func(t *testing.T) { // Setup the client with the test case state c := &vertexAiClient{ - reasoningEngine: tt.existingEngineID, + agentEngineData: &vertexai.AgentEngineData{ + ReasoningEngine: tt.existingEngineID, + }, } // Execute @@ -100,3 +241,240 @@ func TestGetReasoningEngineID(t *testing.T) { }) } } + +func TestAiplatformToGenaiContentPreservesFunctionIDs(t *testing.T) { + args, err := structpb.NewStruct(map[string]any{"city": "Stockholm"}) + if err != nil { + t.Fatalf("structpb.NewStruct(args) failed: %v", err) + } + response, err := structpb.NewStruct(map[string]any{"temperature": 21}) + if err != nil { + t.Fatalf("structpb.NewStruct(response) failed: %v", err) + } + + content := aiplatformToGenaiContent(&aiplatformpb.SessionEvent{ + Content: &aiplatformpb.Content{ + Role: string(genai.RoleModel), + Parts: []*aiplatformpb.Part{ + { + Data: &aiplatformpb.Part_FunctionCall{ + FunctionCall: &aiplatformpb.FunctionCall{ + Id: "call-123", + Name: "get_weather", + Args: args, + }, + }, + }, + { + Data: &aiplatformpb.Part_FunctionResponse{ + FunctionResponse: &aiplatformpb.FunctionResponse{ + Id: "call-123", + Name: "get_weather", + Response: response, + }, + }, + }, + }, + }, + }) + + if content == nil { + t.Fatal("aiplatformToGenaiContent() returned nil content") + } + if got, want := len(content.Parts), 2; got != want { + t.Fatalf("len(content.Parts) = %d, want %d", got, want) + } + + functionCall := content.Parts[0].FunctionCall + if functionCall == nil { + t.Fatal("content.Parts[0].FunctionCall is nil") + } + if got, want := functionCall.ID, "call-123"; got != want { + t.Errorf("FunctionCall.ID = %q, want %q", got, want) + } + + functionResponse := content.Parts[1].FunctionResponse + if functionResponse == nil { + t.Fatal("content.Parts[1].FunctionResponse is nil") + } + if got, want := functionResponse.ID, "call-123"; got != want { + t.Errorf("FunctionResponse.ID = %q, want %q", got, want) + } +} + +func TestToStructPB(t *testing.T) { + tests := []struct { + name string + input any + expectError bool + validate func(t *testing.T, s *structpb.Struct) + }{ + { + name: "simple map representing function call args or function response", + input: map[string]any{"city": "Stockholm"}, + expectError: false, + validate: func(t *testing.T, s *structpb.Struct) { + if got, want := s.Fields["city"].GetStringValue(), "Stockholm"; got != want { + t.Errorf("city = %q, want %q", got, want) + } + }, + }, + { + name: "invalid input", + input: "hello", + expectError: true, + }, + { + name: "custom struct representing possible state delta", + input: struct { + StringValue string + BoolValue bool + IntValue int32 + ArrayValue []string + }{ + StringValue: "value", + BoolValue: false, + IntValue: 123, + ArrayValue: []string{"value"}, + }, + expectError: false, + validate: func(t *testing.T, s *structpb.Struct) { + if _, exists := s.Fields["StringValue"]; !exists { + t.Error("expected 'StringValue' field to exist") + } + if _, exists := s.Fields["BoolValue"]; !exists { + t.Error("expected 'Boolvalue' field to exist") + } + if _, exists := s.Fields["IntValue"]; !exists { + t.Error("expected 'IntValue' field to exist") + } + if _, exists := s.Fields["ArrayValue"]; !exists { + t.Error("expected 'ArrayValue' field to exist") + } + }, + }, + { + name: "custom struct representing possible state delta respects json tags and omitempty", + input: struct { + StringValue string `json:"string_value"` + BoolValue bool `json:"bool_value"` + IntValue int32 `json:"int_value"` + ArrayValue []string `json:"array_value"` + EmptyStringValue string `json:"empty_string_value,omitempty"` + }{ + StringValue: "value", + BoolValue: false, + IntValue: 123, + ArrayValue: []string{"value"}, + EmptyStringValue: "", + }, + expectError: false, + validate: func(t *testing.T, s *structpb.Struct) { + if _, exists := s.Fields["string_value"]; !exists { + t.Error("expected 'string_value' field to exist") + } + if _, exists := s.Fields["bool_value"]; !exists { + t.Error("expected 'bool_value' field to exist") + } + if _, exists := s.Fields["int_value"]; !exists { + t.Error("expected 'int_value' field to exist") + } + if _, exists := s.Fields["array_value"]; !exists { + t.Error("expected 'array_value' field to exist") + } + if _, exists := s.Fields["empty_string_value"]; exists { + t.Error("unexpected 'empty_string_value' field") + } + }, + }, + } + + for _, tt := range tests { + t.Run(tt.name, func(t *testing.T) { + result, err := toStructPB(tt.input) + if (err != nil) != tt.expectError { + t.Errorf("toStructPB() error = %v, expectError %v", err, tt.expectError) + } + if !tt.expectError && tt.validate != nil { + tt.validate(t, result) + } + }) + } +} + +func TestCreateAiplatformpbContent(t *testing.T) { + tests := []struct { + name string + event *session.Event + expectError bool + }{ + { + name: "simple function call args", + event: &session.Event{ + LLMResponse: model.LLMResponse{ + Content: &genai.Content{ + Parts: []*genai.Part{ + genai.NewPartFromFunctionCall("tool", map[string]any{ + "city": "Stockholm", + }), + }, + Role: genai.RoleUser, + }, + }, + }, + expectError: false, + }, + { + name: "simple function response", + event: &session.Event{ + LLMResponse: model.LLMResponse{ + Content: &genai.Content{ + Parts: []*genai.Part{ + genai.NewPartFromFunctionResponse("tool", map[string]any{ + "city": "Stockholm", + }), + }, + Role: genai.RoleUser, + }, + }, + }, + expectError: false, + }, + } + for _, tt := range tests { + t.Run(tt.name, func(t *testing.T) { + _, err := createAiplatformpbContent(tt.event) + if (err != nil) != tt.expectError { + t.Errorf("createAiplatformpbContent() error = %v, expectError %v", err, tt.expectError) + } + }) + } +} + +func TestCreateAiplatformpbMetadata(t *testing.T) { + tests := []struct { + name string + event *session.Event + expectError bool + }{ + { + name: "simple custom metadata", + event: &session.Event{ + LLMResponse: model.LLMResponse{ + CustomMetadata: map[string]any{ + "key": "value", + }, + }, + }, + expectError: false, + }, + } + for _, tt := range tests { + t.Run(tt.name, func(t *testing.T) { + _, err := createAiplatformpbMetadata(tt.event) + if (err != nil) != tt.expectError { + t.Errorf("createAiplatformpbMetadata() error = %v, expectError %v", err, tt.expectError) + } + }) + } +} diff --git a/tool/agenttool/agent_tool.go b/tool/agenttool/agent_tool.go index 342889d16..8041d2bcc 100644 --- a/tool/agenttool/agent_tool.go +++ b/tool/agenttool/agent_tool.go @@ -118,7 +118,7 @@ func (t *agentTool) Declaration() *genai.FunctionDeclaration { // Run executes the wrapped agent with the provided arguments. // It creates a new session for the sub-agent, runs the agent, and returns // the final result. -func (t *agentTool) Run(toolCtx tool.Context, args any) (map[string]any, error) { +func (t *agentTool) Run(toolCtx agent.ToolContext, args any) (map[string]any, error) { margs, ok := args.(map[string]any) if !ok { return nil, fmt.Errorf("agentTool expects map[string]any arguments, got %T", args) @@ -251,6 +251,6 @@ func (t *agentTool) Run(toolCtx tool.Context, args any) (map[string]any, error) } // ProcessRequest adds the agent tool's function declaration to the LLM request. -func (t *agentTool) ProcessRequest(ctx tool.Context, req *model.LLMRequest) error { +func (t *agentTool) ProcessRequest(ctx agent.ToolContext, req *model.LLMRequest) error { return toolutils.PackTool(req, t) } diff --git a/tool/agenttool/agent_tool_test.go b/tool/agenttool/agent_tool_test.go index 5eae37ba6..d37b3e488 100644 --- a/tool/agenttool/agent_tool_test.go +++ b/tool/agenttool/agent_tool_test.go @@ -29,7 +29,6 @@ import ( "google.golang.org/adk/model" "google.golang.org/adk/model/gemini" "google.golang.org/adk/session" - "google.golang.org/adk/tool" "google.golang.org/adk/tool/agenttool" ) @@ -440,7 +439,7 @@ func createAgentWithModel(t *testing.T, inputSchema, outputSchema *genai.Schema, return agent } -func createToolContext(t *testing.T, testAgent agent.Agent) tool.Context { +func createToolContext(t *testing.T, testAgent agent.Agent) agent.ToolContext { t.Helper() sessionService := session.InMemoryService() @@ -457,5 +456,5 @@ func createToolContext(t *testing.T, testAgent agent.Agent) tool.Context { Session: createResponse.Session, }) - return toolinternal.NewToolContext(ctx, "", &session.EventActions{}, nil) + return agent.NewToolContext(ctx, "", &session.EventActions{}, nil) } diff --git a/tool/context.go b/tool/context.go new file mode 100644 index 000000000..6ca0f9bd8 --- /dev/null +++ b/tool/context.go @@ -0,0 +1,30 @@ +// Copyright 2026 Google LLC +// +// Licensed under the Apache License, Version 2.0 (the "License"); +// you may not use this file except in compliance with the License. +// You may obtain a copy of the License at +// +// http://www.apache.org/licenses/LICENSE-2.0 +// +// Unless required by applicable law or agreed to in writing, software +// distributed under the License is distributed on an "AS IS" BASIS, +// WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. +// See the License for the specific language governing permissions and +// limitations under the License. + +package tool + +import ( + "google.golang.org/adk/agent" + "google.golang.org/adk/session" + "google.golang.org/adk/tool/toolconfirmation" +) + +// NewToolContext constructs a ToolContext for a tool execution. +// +// Deprecated: use agent.NewToolContext directly. This wrapper exists only +// to minimize churn during the migration and will be removed in a future +// release. +func NewToolContext(ic agent.InvocationContext, functionCallID string, actions *session.EventActions, confirmation *toolconfirmation.ToolConfirmation) Context { + return agent.NewToolContext(ic, functionCallID, actions, confirmation) +} diff --git a/internal/toolinternal/context_test.go b/tool/context_test.go similarity index 75% rename from internal/toolinternal/context_test.go rename to tool/context_test.go index a846e78e3..2596941ed 100644 --- a/internal/toolinternal/context_test.go +++ b/tool/context_test.go @@ -1,4 +1,4 @@ -// Copyright 2025 Google LLC +// Copyright 2026 Google LLC // // Licensed under the Apache License, Version 2.0 (the "License"); // you may not use this file except in compliance with the License. @@ -12,19 +12,19 @@ // See the License for the specific language governing permissions and // limitations under the License. -package toolinternal +package tool_test import ( "testing" "google.golang.org/adk/agent" - contextinternal "google.golang.org/adk/internal/context" + icontext "google.golang.org/adk/internal/context" "google.golang.org/adk/session" ) -func TestToolContext(t *testing.T) { - inv := contextinternal.NewInvocationContext(t.Context(), contextinternal.InvocationContextParams{}) - toolCtx := NewToolContext(inv, "fn1", &session.EventActions{}, nil) +func TestNewToolContext_Interfaces(t *testing.T) { + inv := icontext.NewInvocationContext(t.Context(), icontext.InvocationContextParams{}) + toolCtx := agent.NewToolContext(inv, "fn1", &session.EventActions{}, nil) if _, ok := toolCtx.(agent.ReadonlyContext); !ok { t.Errorf("ToolContext(%+T) is unexpectedly not a ReadonlyContext", toolCtx) @@ -34,10 +34,10 @@ func TestToolContext(t *testing.T) { } } -func TestRequestConfirmation_SetsSkipSummarization(t *testing.T) { - inv := contextinternal.NewInvocationContext(t.Context(), contextinternal.InvocationContextParams{}) +func TestNewToolContext_RequestConfirmation_SetsSkipSummarization(t *testing.T) { + inv := icontext.NewInvocationContext(t.Context(), icontext.InvocationContextParams{}) actions := &session.EventActions{} - toolCtx := NewToolContext(inv, "fn1", actions, nil) + toolCtx := agent.NewToolContext(inv, "fn1", actions, nil) err := toolCtx.RequestConfirmation("please confirm", map[string]any{"key": "value"}) if err != nil { @@ -63,11 +63,11 @@ func TestRequestConfirmation_SetsSkipSummarization(t *testing.T) { } } -func TestRequestConfirmation_AutoGeneratesIDWhenEmpty(t *testing.T) { - inv := contextinternal.NewInvocationContext(t.Context(), contextinternal.InvocationContextParams{}) +func TestNewToolContext_RequestConfirmation_AutoGeneratesIDWhenEmpty(t *testing.T) { + inv := icontext.NewInvocationContext(t.Context(), icontext.InvocationContextParams{}) actions := &session.EventActions{} // NewToolContext auto-generates a UUID when functionCallID is empty. - toolCtx := NewToolContext(inv, "", actions, nil) + toolCtx := agent.NewToolContext(inv, "", actions, nil) err := toolCtx.RequestConfirmation("hint", nil) if err != nil { diff --git a/tool/exampletool/tool.go b/tool/exampletool/tool.go index cb67382e3..4592fe6f0 100644 --- a/tool/exampletool/tool.go +++ b/tool/exampletool/tool.go @@ -21,9 +21,9 @@ import ( "google.golang.org/genai" + "google.golang.org/adk/agent" "google.golang.org/adk/internal/utils" "google.golang.org/adk/model" - "google.golang.org/adk/tool" ) type Example struct { @@ -55,7 +55,7 @@ func (s exampleTool) Description() string { } // ProcessRequest adds the exampleTool examples to the LLM request. -func (s exampleTool) ProcessRequest(ctx tool.Context, req *model.LLMRequest) error { +func (s exampleTool) ProcessRequest(ctx agent.ToolContext, req *model.LLMRequest) error { parts := ctx.UserContent().Parts if len(parts) == 0 || parts[0].Text == "" { return nil diff --git a/tool/exitlooptool/tool.go b/tool/exitlooptool/tool.go index bb42ad2a7..aea1d5c6e 100644 --- a/tool/exitlooptool/tool.go +++ b/tool/exitlooptool/tool.go @@ -18,11 +18,12 @@ package exitlooptool import ( "fmt" + "google.golang.org/adk/agent" "google.golang.org/adk/tool" "google.golang.org/adk/tool/functiontool" ) -func exitLoop(ctx tool.Context, myArgs struct{}) (map[string]string, error) { +func exitLoop(ctx agent.ToolContext, myArgs struct{}) (map[string]string, error) { ctx.Actions().Escalate = true ctx.Actions().SkipSummarization = true return map[string]string{}, nil diff --git a/tool/functiontool/function.go b/tool/functiontool/function.go index e3367dd45..54e8ef055 100644 --- a/tool/functiontool/function.go +++ b/tool/functiontool/function.go @@ -24,6 +24,7 @@ import ( "github.com/google/jsonschema-go/jsonschema" "google.golang.org/genai" + "google.golang.org/adk/agent" "google.golang.org/adk/internal/toolinternal/toolutils" "google.golang.org/adk/internal/typeutil" "google.golang.org/adk/model" @@ -67,8 +68,8 @@ type Config struct { } // Func represents a Go function that can be wrapped in a tool. -// It takes a tool.Context and a generic argument type, and returns a generic result type. -type Func[TArgs, TResults any] func(tool.Context, TArgs) (TResults, error) +// It takes a agent.ToolContext and a generic argument type, and returns a generic result type. +type Func[TArgs, TResults any] func(agent.ToolContext, TArgs) (TResults, error) // ErrInvalidArgument indicates the input parameter type is invalid. var ErrInvalidArgument = errors.New("invalid argument") @@ -151,7 +152,7 @@ func (f *functionTool[TArgs, TResults]) IsLongRunning() bool { } // ProcessRequest packs the function tool's declaration into the LLM request. -func (f *functionTool[TArgs, TResults]) ProcessRequest(ctx tool.Context, req *model.LLMRequest) error { +func (f *functionTool[TArgs, TResults]) ProcessRequest(ctx agent.ToolContext, req *model.LLMRequest) error { return toolutils.PackTool(req, f) } @@ -181,7 +182,7 @@ func (f *functionTool[TArgs, TResults]) Declaration() *genai.FunctionDeclaration } // Run executes the tool with the provided context and yields events. -func (f *functionTool[TArgs, TResults]) Run(ctx tool.Context, args any) (result map[string]any, err error) { +func (f *functionTool[TArgs, TResults]) Run(ctx agent.ToolContext, args any) (result map[string]any, err error) { // TODO: Handle function call request from tc.InvocationContext. defer func() { if r := recover(); r != nil { diff --git a/tool/functiontool/function_test.go b/tool/functiontool/function_test.go index 570b726ab..053083334 100644 --- a/tool/functiontool/function_test.go +++ b/tool/functiontool/function_test.go @@ -28,6 +28,7 @@ import ( "github.com/google/jsonschema-go/jsonschema" "google.golang.org/genai" + "google.golang.org/adk/agent" "google.golang.org/adk/agent/llmagent" icontext "google.golang.org/adk/internal/context" "google.golang.org/adk/internal/testutil" @@ -49,7 +50,7 @@ type SumResult struct { Sum int `json:"sum"` // the sum of two integers } -func sumFunc(ctx tool.Context, input SumArgs) (SumResult, error) { +func sumFunc(ctx agent.ToolContext, input SumArgs) (SumResult, error) { return SumResult{Sum: input.A + input.B}, nil } @@ -64,9 +65,9 @@ func ExampleNew() { _ = sumTool // use the tool } -func createToolContext(t *testing.T) tool.Context { +func createToolContext(t *testing.T) agent.ToolContext { invCtx := icontext.NewInvocationContext(t.Context(), icontext.InvocationContextParams{}) - return toolinternal.NewToolContext(invCtx, "", &session.EventActions{}, nil) + return agent.NewToolContext(invCtx, "", &session.EventActions{}, nil) } //go:generate go test -v -httprecord=.* @@ -100,7 +101,7 @@ func TestFunctionTool_Simple(t *testing.T) { }, } - weatherReport := func(ctx tool.Context, input Args) (Result, error) { + weatherReport := func(ctx agent.ToolContext, input Args) (Result, error) { city := strings.ToLower(input.City) if ret, ok := resultSet[city]; ok { return ret, nil @@ -201,7 +202,7 @@ func TestFunctionTool_DifferentFunctionDeclarations_ConsolidatedInOneGenAiTool(t type IntOutput struct { Result int `json:"result"` } - identityFunc := func(ctx tool.Context, input IntInput) (IntOutput, error) { + identityFunc := func(ctx agent.ToolContext, input IntInput) (IntOutput, error) { return IntOutput{Result: input.X}, nil } identityTool, err := functiontool.New(functiontool.Config{ @@ -219,7 +220,7 @@ func TestFunctionTool_DifferentFunctionDeclarations_ConsolidatedInOneGenAiTool(t type StringOutput struct { Result string `json:"result"` } - stringIdentityFunc := func(ctx tool.Context, input StringInput) (StringOutput, error) { + stringIdentityFunc := func(ctx agent.ToolContext, input StringInput) (StringOutput, error) { return StringOutput{Result: input.Value}, nil } stringIdentityTool, err := functiontool.New( @@ -265,7 +266,7 @@ func TestFunctionTool_ReturnsBasicType(t *testing.T) { "paris": "The weather in Paris is sunny with a temperature of 25 derees Celsius.", } - weatherReport := func(ctx tool.Context, input Args) (string, error) { + weatherReport := func(ctx agent.ToolContext, input Args) (string, error) { city := strings.ToLower(input.City) if ret, ok := resultSet[city]; ok { return ret, nil @@ -355,7 +356,7 @@ func TestFunctionTool_MapInput(t *testing.T) { Name: "sum_map", Description: "sums numbers provided in a map input", }, - func(ctx tool.Context, input map[string]int) (Output, error) { + func(ctx agent.ToolContext, input map[string]int) (Output, error) { return Output{Sum: input["a"] + input["b"]}, nil }) if err != nil { @@ -440,7 +441,7 @@ func TestFunctionTool_CustomSchema(t *testing.T) { Name: "print_quantity", Description: "print the remaining quantity of the given fruit.", InputSchema: ischema, - }, func(ctx tool.Context, input Args) (any, error) { + }, func(ctx agent.ToolContext, input Args) (any, error) { fruit := strings.ToLower(input.Fruit) if fruit != "mandarin" && fruit != "kiwi" { t.Errorf("unexpected fruit: %q", fruit) @@ -552,7 +553,7 @@ type SimpleArgs struct { Num int } -func okFunc(_ tool.Context, _ SimpleArgs) (string, error) { +func okFunc(_ agent.ToolContext, _ SimpleArgs) (string, error) { return "ok", nil } @@ -584,6 +585,9 @@ func TestToolConfirmation(t *testing.T) { args: map[string]any{"Num": 1}, want: []*genai.Content{ genai.NewContentFromFunctionCall("test_tool", map[string]any{"Num": 1}, "model"), + genai.NewContentFromFunctionResponse("test_tool", map[string]any{ + "error": errors.New("error tool \"test_tool\" requires confirmation, please approve or reject"), + }, "user"), genai.NewContentFromFunctionCall(toolconfirmation.FunctionCallName, map[string]any{ "originalFunctionCall": &genai.FunctionCall{ Args: map[string]any{"Num": 1}, @@ -593,9 +597,6 @@ func TestToolConfirmation(t *testing.T) { Hint: "Please approve or reject the tool call test_tool() by responding with a FunctionResponse with an expected ToolConfirmation payload.", }, }, "model"), - genai.NewContentFromFunctionResponse("test_tool", map[string]any{ - "error": errors.New("error tool \"test_tool\" requires confirmation, please approve or reject"), - }, "user"), }, }, { @@ -608,6 +609,9 @@ func TestToolConfirmation(t *testing.T) { confirmFunctionResponse: &genai.FunctionResponse{Name: toolconfirmation.FunctionCallName, Response: map[string]any{"confirmed": true}}, want: []*genai.Content{ genai.NewContentFromFunctionCall("test_tool", map[string]any{"Num": 1}, "model"), + genai.NewContentFromFunctionResponse("test_tool", map[string]any{ + "error": errors.New("error tool \"test_tool\" requires confirmation, please approve or reject"), + }, "user"), genai.NewContentFromFunctionCall(toolconfirmation.FunctionCallName, map[string]any{ "originalFunctionCall": &genai.FunctionCall{ Args: map[string]any{"Num": 1}, @@ -617,9 +621,6 @@ func TestToolConfirmation(t *testing.T) { Hint: "Please approve or reject the tool call test_tool() by responding with a FunctionResponse with an expected ToolConfirmation payload.", }, }, "model"), - genai.NewContentFromFunctionResponse("test_tool", map[string]any{ - "error": errors.New("error tool \"test_tool\" requires confirmation, please approve or reject"), - }, "user"), genai.NewContentFromFunctionResponse("test_tool", map[string]any{"result": "ok"}, "user"), }, }, @@ -633,6 +634,9 @@ func TestToolConfirmation(t *testing.T) { confirmFunctionResponse: &genai.FunctionResponse{Name: toolconfirmation.FunctionCallName, Response: map[string]any{"confirmed": false}}, want: []*genai.Content{ genai.NewContentFromFunctionCall("test_tool", map[string]any{"Num": 1}, "model"), + genai.NewContentFromFunctionResponse("test_tool", map[string]any{ + "error": errors.New("error tool \"test_tool\" requires confirmation, please approve or reject"), + }, "user"), genai.NewContentFromFunctionCall(toolconfirmation.FunctionCallName, map[string]any{ "originalFunctionCall": &genai.FunctionCall{ Args: map[string]any{"Num": 1}, @@ -642,9 +646,6 @@ func TestToolConfirmation(t *testing.T) { Hint: "Please approve or reject the tool call test_tool() by responding with a FunctionResponse with an expected ToolConfirmation payload.", }, }, "model"), - genai.NewContentFromFunctionResponse("test_tool", map[string]any{ - "error": errors.New("error tool \"test_tool\" requires confirmation, please approve or reject"), - }, "user"), genai.NewContentFromFunctionResponse("test_tool", map[string]any{ "error": errors.New("error tool \"test_tool\" call is rejected"), }, "user"), @@ -675,6 +676,9 @@ func TestToolConfirmation(t *testing.T) { args: map[string]any{"Num": 4}, want: []*genai.Content{ genai.NewContentFromFunctionCall("test_tool", map[string]any{"Num": 4}, "model"), + genai.NewContentFromFunctionResponse("test_tool", map[string]any{ + "error": errors.New("error tool \"test_tool\" requires confirmation, please approve or reject"), + }, "user"), genai.NewContentFromFunctionCall(toolconfirmation.FunctionCallName, map[string]any{ "originalFunctionCall": &genai.FunctionCall{ Args: map[string]any{"Num": 4}, @@ -684,9 +688,6 @@ func TestToolConfirmation(t *testing.T) { Hint: "Please approve or reject the tool call test_tool() by responding with a FunctionResponse with an expected ToolConfirmation payload.", }, }, "model"), - genai.NewContentFromFunctionResponse("test_tool", map[string]any{ - "error": errors.New("error tool \"test_tool\" requires confirmation, please approve or reject"), - }, "user"), }, }, { @@ -701,6 +702,9 @@ func TestToolConfirmation(t *testing.T) { confirmFunctionResponse: &genai.FunctionResponse{Name: toolconfirmation.FunctionCallName, Response: map[string]any{"confirmed": true}}, want: []*genai.Content{ genai.NewContentFromFunctionCall("test_tool", map[string]any{"Num": 4}, "model"), + genai.NewContentFromFunctionResponse("test_tool", map[string]any{ + "error": errors.New("error tool \"test_tool\" requires confirmation, please approve or reject"), + }, "user"), genai.NewContentFromFunctionCall(toolconfirmation.FunctionCallName, map[string]any{ "originalFunctionCall": &genai.FunctionCall{ Args: map[string]any{"Num": 4}, @@ -710,9 +714,6 @@ func TestToolConfirmation(t *testing.T) { Hint: "Please approve or reject the tool call test_tool() by responding with a FunctionResponse with an expected ToolConfirmation payload.", }, }, "model"), - genai.NewContentFromFunctionResponse("test_tool", map[string]any{ - "error": errors.New("error tool \"test_tool\" requires confirmation, please approve or reject"), - }, "user"), genai.NewContentFromFunctionResponse("test_tool", map[string]any{"result": "ok"}, "user"), }, }, @@ -728,6 +729,9 @@ func TestToolConfirmation(t *testing.T) { confirmFunctionResponse: &genai.FunctionResponse{Name: toolconfirmation.FunctionCallName, Response: map[string]any{"confirmed": false}}, want: []*genai.Content{ genai.NewContentFromFunctionCall("test_tool", map[string]any{"Num": 4}, "model"), + genai.NewContentFromFunctionResponse("test_tool", map[string]any{ + "error": errors.New("error tool \"test_tool\" requires confirmation, please approve or reject"), + }, "user"), genai.NewContentFromFunctionCall(toolconfirmation.FunctionCallName, map[string]any{ "originalFunctionCall": &genai.FunctionCall{ Args: map[string]any{"Num": 4}, @@ -737,9 +741,6 @@ func TestToolConfirmation(t *testing.T) { Hint: "Please approve or reject the tool call test_tool() by responding with a FunctionResponse with an expected ToolConfirmation payload.", }, }, "model"), - genai.NewContentFromFunctionResponse("test_tool", map[string]any{ - "error": errors.New("error tool \"test_tool\" requires confirmation, please approve or reject"), - }, "user"), genai.NewContentFromFunctionResponse("test_tool", map[string]any{ "error": errors.New("error tool \"test_tool\" call is rejected"), }, "user"), @@ -865,6 +866,56 @@ func TestToolConfirmation(t *testing.T) { } } +func TestToolConfirmationYieldsFunctionResponseBeforeConfirmation(t *testing.T) { + mockModel := &testutil.MockModel{ + Responses: []*genai.Content{ + genai.NewContentFromFunctionCall("test_tool", map[string]any{"Num": 1}, genai.RoleModel), + }, + } + + myTool, err := functiontool.New(functiontool.Config{ + Name: "test_tool", + RequireConfirmation: true, + }, okFunc) + if err != nil { + t.Fatalf("Failed to create tool: %v", err) + } + + a, err := llmagent.New(llmagent.Config{ + Name: "simple agent", + Model: mockModel, + Tools: []tool.Tool{myTool}, + }) + if err != nil { + t.Fatalf("failed to create llm agent: %v", err) + } + + runner := testutil.NewTestAgentRunner(t, a) + sawFunctionResponse := false + + for got, err := range runner.Run(t, "id", "message") { + // The mock model emits a sentinel error once exhausted; treat any + // error as end-of-stream and rely on the post-loop assertions. + if err != nil { + break + } + + for _, part := range got.Content.Parts { + if part.FunctionResponse != nil && part.FunctionResponse.Name == "test_tool" { + sawFunctionResponse = true + } + if part.FunctionCall != nil && part.FunctionCall.Name == toolconfirmation.FunctionCallName { + if !sawFunctionResponse { + t.Fatal("confirmation was yielded before the tool function response") + } + return + } + } + } + + t.Fatal("expected confirmation event") +} + // Mock types for TArgs and TResults type TestArgs struct { Name string @@ -876,7 +927,7 @@ type TestResult struct { func TestNew_RequireConfirmationProvider_Validation(t *testing.T) { // A dummy handler to satisfy the function signature - dummyHandler := func(_ tool.Context, _ TestArgs) (TestResult, error) { + dummyHandler := func(_ agent.ToolContext, _ TestArgs) (TestResult, error) { return TestResult{Value: 1}, nil } @@ -987,7 +1038,7 @@ func TestNew_InvalidInputType(t *testing.T) { return functiontool.New(functiontool.Config{ Name: "string_tool", Description: "a tool with string input", - }, func(ctx tool.Context, input string) (string, error) { + }, func(ctx agent.ToolContext, input string) (string, error) { return input, nil }) }, @@ -999,7 +1050,7 @@ func TestNew_InvalidInputType(t *testing.T) { return functiontool.New(functiontool.Config{ Name: "int_tool", Description: "a tool with int input", - }, func(ctx tool.Context, input int) (int, error) { + }, func(ctx agent.ToolContext, input int) (int, error) { return input, nil }) }, @@ -1011,7 +1062,7 @@ func TestNew_InvalidInputType(t *testing.T) { return functiontool.New(functiontool.Config{ Name: "bool_tool", Description: "a tool with bool input", - }, func(ctx tool.Context, input bool) (bool, error) { + }, func(ctx agent.ToolContext, input bool) (bool, error) { return input, nil }) }, @@ -1037,7 +1088,7 @@ func TestFunctionTool_PanicRecovery(t *testing.T) { Value string `json:"value"` } - panicHandler := func(ctx tool.Context, input Args) (string, error) { + panicHandler := func(ctx agent.ToolContext, input Args) (string, error) { panic("intentional panic for testing") } diff --git a/tool/functiontool/long_running_function_test.go b/tool/functiontool/long_running_function_test.go index 0868c7cfd..d81510801 100644 --- a/tool/functiontool/long_running_function_test.go +++ b/tool/functiontool/long_running_function_test.go @@ -23,6 +23,7 @@ import ( "github.com/google/go-cmp/cmp/cmpopts" "google.golang.org/genai" + "google.golang.org/adk/agent" "google.golang.org/adk/agent/llmagent" "google.golang.org/adk/internal/testutil" "google.golang.org/adk/internal/toolinternal" @@ -39,7 +40,7 @@ func TestNewLongRunningFunctionTool(t *testing.T) { Result string `json:"result"` // the operation result } - handler := func(ctx tool.Context, input SumArgs) (SumResult, error) { + handler := func(ctx agent.ToolContext, input SumArgs) (SumResult, error) { return SumResult{Result: "Processing sum"}, nil } sumTool, err := functiontool.New(functiontool.Config{ @@ -80,7 +81,7 @@ type IncArgs struct{} func TestLongRunningFunctionFlow(t *testing.T) { functionCalled := 0 - increaseByOne := func(ctx tool.Context, x IncArgs) (map[string]string, error) { + increaseByOne := func(ctx agent.ToolContext, x IncArgs) (map[string]string, error) { functionCalled++ return map[string]string{"status": "pending"}, nil } @@ -89,7 +90,7 @@ func TestLongRunningFunctionFlow(t *testing.T) { func TestLongRunningStringFunctionFlow(t *testing.T) { functionCalled := 0 - increaseByOne := func(ctx tool.Context, x IncArgs) (string, error) { + increaseByOne := func(ctx agent.ToolContext, x IncArgs) (string, error) { functionCalled++ return "pending", nil } @@ -97,7 +98,7 @@ func TestLongRunningStringFunctionFlow(t *testing.T) { } // --- Test Suite --- -func testLongRunningFunctionFlow[Out any](t *testing.T, increaseByOne func(ctx tool.Context, x IncArgs) (Out, error), resultKey string, callCount *int) { +func testLongRunningFunctionFlow[Out any](t *testing.T, increaseByOne func(ctx agent.ToolContext, x IncArgs) (Out, error), resultKey string, callCount *int) { // 1. Setup responses := []*genai.Content{ genai.NewContentFromFunctionCall("increaseByOne", map[string]any{}, "model"), @@ -266,7 +267,7 @@ func TestLongRunningToolIDsAreSet(t *testing.T) { type IncArgs struct{} - increaseByOne := func(ctx tool.Context, x IncArgs) (map[string]string, error) { + increaseByOne := func(ctx agent.ToolContext, x IncArgs) (map[string]string, error) { functionCalled++ return map[string]string{"status": "pending"}, nil } diff --git a/tool/functiontool/streaming_function.go b/tool/functiontool/streaming_function.go new file mode 100644 index 000000000..af3d16971 --- /dev/null +++ b/tool/functiontool/streaming_function.go @@ -0,0 +1,181 @@ +// Copyright 2025 Google LLC +// +// Licensed under the Apache License, Version 2.0 (the "License"); +// you may not use this file except in compliance with the License. +// You may obtain a copy of the License at +// +// http://www.apache.org/licenses/LICENSE-2.0 +// +// Unless required by applicable law or agreed to in writing, software +// distributed under the License is distributed on an "AS IS" BASIS, +// WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. +// See the License for the specific language governing permissions and +// limitations under the License. + +// Package functiontool provides a tool that wraps a Go function. +package functiontool + +import ( + "fmt" + "iter" + "reflect" + "runtime/debug" + + "github.com/google/jsonschema-go/jsonschema" + "google.golang.org/genai" + + "google.golang.org/adk/agent" + "google.golang.org/adk/internal/toolinternal/toolutils" + "google.golang.org/adk/internal/typeutil" + "google.golang.org/adk/model" + "google.golang.org/adk/tool" +) + +// StreamingFunc represents a Go function that streams results. +type StreamingFunc[TArgs any] func(agent.ToolContext, TArgs) iter.Seq2[string, error] + +// NewStreaming creates a new streaming tool. +func NewStreaming[TArgs any](cfg Config, handler StreamingFunc[TArgs]) (tool.Tool, error) { + var zeroArgs TArgs + argsType := reflect.TypeOf(zeroArgs) + for argsType != nil && argsType.Kind() == reflect.Pointer { + argsType = argsType.Elem() + } + if argsType == nil || (argsType.Kind() != reflect.Struct && argsType.Kind() != reflect.Map) { + return nil, fmt.Errorf("input must be a struct or a map or a pointer to those types, but received: %v: %w", argsType, ErrInvalidArgument) + } + + ischema, err := resolvedSchema[TArgs](cfg.InputSchema) + if err != nil { + return nil, fmt.Errorf("failed to infer input schema: %w", err) + } + + var confirmWrapper func(TArgs) bool + + if cfg.RequireConfirmationProvider != nil { + fn, ok := cfg.RequireConfirmationProvider.(func(TArgs) bool) + if !ok { + return nil, fmt.Errorf("error RequireConfirmationProvider must be a function with signature func(%T) bool", *new(TArgs)) + } + confirmWrapper = fn + } + + return &streamingFunctionTool[TArgs]{ + cfg: cfg, + inputSchema: ischema, + handler: handler, + requireConfirmation: cfg.RequireConfirmation, + requireConfirmationProvider: confirmWrapper, + }, nil +} + +// streamingFunctionTool wraps a Go function that streams results. +type streamingFunctionTool[TArgs any] struct { + cfg Config + + // A JSON Schema object defining the expected parameters for the tool. + inputSchema *jsonschema.Resolved + + // handler is the Go function. + handler StreamingFunc[TArgs] + + requireConfirmation bool + + requireConfirmationProvider func(TArgs) bool +} + +// Description implements tool.Tool. +func (f *streamingFunctionTool[TArgs]) Description() string { + return f.cfg.Description +} + +// Name implements tool.Tool. +func (f *streamingFunctionTool[TArgs]) Name() string { + return f.cfg.Name +} + +// IsLongRunning implements tool.Tool. +func (f *streamingFunctionTool[TArgs]) IsLongRunning() bool { + return f.cfg.IsLongRunning +} + +// ProcessRequest packs the function tool's declaration into the LLM request. +func (f *streamingFunctionTool[TArgs]) ProcessRequest(ctx agent.ToolContext, req *model.LLMRequest) error { + return toolutils.PackTool(req, f) +} + +// FunctionDeclaration implements toolinternal.StreamingFunctionTool. +func (f *streamingFunctionTool[TArgs]) Declaration() *genai.FunctionDeclaration { + decl := &genai.FunctionDeclaration{ + Name: f.Name(), + Description: f.Description(), + } + if f.inputSchema != nil { + decl.ParametersJsonSchema = f.inputSchema.Schema() + } + + if f.cfg.IsLongRunning { + instruction := "NOTE: This is a long-running operation. Do not call this tool again if it has already returned some intermediate or pending status." + if decl.Description != "" { + decl.Description += "\n\n" + instruction + } else { + decl.Description = instruction + } + } + + return decl +} + +// RunStream executes the tool with the provided context and yields events. +func (f *streamingFunctionTool[TArgs]) RunStream(ctx agent.ToolContext, args any) iter.Seq2[string, error] { + return func(yield func(string, error) bool) { + defer func() { + if r := recover(); r != nil { + yield("", fmt.Errorf("panic in tool %q: %v\nstack: %s", f.Name(), r, debug.Stack())) + } + }() + + m, ok := args.(map[string]any) + if !ok { + yield("", fmt.Errorf("unexpected args type, got: %T", args)) + return + } + input, err := typeutil.ConvertToWithJSONSchema[map[string]any, TArgs](m, f.inputSchema) + if err != nil { + yield("", err) + return + } + + if confirmation := ctx.ToolConfirmation(); confirmation != nil { + if !confirmation.Confirmed { + yield("", fmt.Errorf("error tool %q %w", f.Name(), tool.ErrConfirmationRejected)) + return + } + } else { + requireConfirmation := f.requireConfirmation + + if f.requireConfirmationProvider != nil { + requireConfirmation = f.requireConfirmationProvider(input) + } + + if requireConfirmation { + err := ctx.RequestConfirmation( + fmt.Sprintf("Please approve or reject the tool call %s() by responding with a FunctionResponse with an expected ToolConfirmation payload.", + f.Name()), nil) + if err != nil { + yield("", err) + return + } + ctx.Actions().SkipSummarization = true + yield("", fmt.Errorf("error tool %q %w", f.Name(), tool.ErrConfirmationRequired)) + return + } + } + + for res, err := range f.handler(ctx, input) { + if !yield(res, err) { + return + } + } + } +} diff --git a/tool/geminitool/google_search.go b/tool/geminitool/google_search.go index 7002726fd..74bd3c3de 100644 --- a/tool/geminitool/google_search.go +++ b/tool/geminitool/google_search.go @@ -17,8 +17,8 @@ package geminitool import ( "google.golang.org/genai" + "google.golang.org/adk/agent" "google.golang.org/adk/model" - "google.golang.org/adk/tool" ) // GoogleSearch is a built-in tool that is automatically invoked by Gemini 2 @@ -38,7 +38,7 @@ func (s GoogleSearch) Description() string { } // ProcessRequest adds the GoogleSearch tool to the LLM request. -func (s GoogleSearch) ProcessRequest(ctx tool.Context, req *model.LLMRequest) error { +func (s GoogleSearch) ProcessRequest(ctx agent.ToolContext, req *model.LLMRequest) error { return setTool(req, &genai.Tool{ GoogleSearch: &genai.GoogleSearch{}, }) diff --git a/tool/geminitool/tool.go b/tool/geminitool/tool.go index 37e8199a5..117e01206 100644 --- a/tool/geminitool/tool.go +++ b/tool/geminitool/tool.go @@ -34,6 +34,7 @@ import ( "google.golang.org/genai" + "google.golang.org/adk/agent" "google.golang.org/adk/model" "google.golang.org/adk/tool" ) @@ -55,7 +56,7 @@ type geminiTool struct { } // ProcessRequest adds the Gemini tool to the LLM request. -func (t *geminiTool) ProcessRequest(ctx tool.Context, req *model.LLMRequest) error { +func (t *geminiTool) ProcessRequest(ctx agent.ToolContext, req *model.LLMRequest) error { return setTool(req, t.value) } diff --git a/tool/loadartifactstool/load_artifacts_tool.go b/tool/loadartifactstool/load_artifacts_tool.go index 80b034777..b912926ea 100644 --- a/tool/loadartifactstool/load_artifacts_tool.go +++ b/tool/loadartifactstool/load_artifacts_tool.go @@ -85,7 +85,7 @@ func (t *artifactsTool) Declaration() *genai.FunctionDeclaration { } // Run implements tool.Tool. -func (t *artifactsTool) Run(ctx tool.Context, args any) (map[string]any, error) { +func (t *artifactsTool) Run(ctx agent.ToolContext, args any) (map[string]any, error) { m, ok := args.(map[string]any) if !ok { return nil, fmt.Errorf("unexpected args type, got: %T", args) @@ -117,7 +117,7 @@ func (t *artifactsTool) Run(ctx tool.Context, args any) (map[string]any, error) // ProcessRequest processes the LLM request. It packs the tool, appends initial // instructions, and processes any load artifacts function calls. -func (t *artifactsTool) ProcessRequest(ctx tool.Context, req *model.LLMRequest) error { +func (t *artifactsTool) ProcessRequest(ctx agent.ToolContext, req *model.LLMRequest) error { if err := toolutils.PackTool(req, t); err != nil { return err } @@ -127,7 +127,7 @@ func (t *artifactsTool) ProcessRequest(ctx tool.Context, req *model.LLMRequest) return t.processLoadArtifactsFunctionCall(ctx, req) } -func (t *artifactsTool) appendInitialInstructions(ctx tool.Context, req *model.LLMRequest) error { +func (t *artifactsTool) appendInitialInstructions(ctx agent.ToolContext, req *model.LLMRequest) error { resp, err := ctx.Artifacts().List(ctx) if err != nil { return fmt.Errorf("failed to list artifacts: %w", err) @@ -151,7 +151,7 @@ func (t *artifactsTool) appendInitialInstructions(ctx tool.Context, req *model.L return nil } -func (t *artifactsTool) processLoadArtifactsFunctionCall(ctx tool.Context, req *model.LLMRequest) error { +func (t *artifactsTool) processLoadArtifactsFunctionCall(ctx agent.ToolContext, req *model.LLMRequest) error { if len(req.Contents) == 0 { return nil } diff --git a/tool/loadartifactstool/load_artifacts_tool_test.go b/tool/loadartifactstool/load_artifacts_tool_test.go index 93366776a..846fa0e70 100644 --- a/tool/loadartifactstool/load_artifacts_tool_test.go +++ b/tool/loadartifactstool/load_artifacts_tool_test.go @@ -21,12 +21,12 @@ import ( "github.com/google/go-cmp/cmp" "google.golang.org/genai" + "google.golang.org/adk/agent" "google.golang.org/adk/artifact" artifactinternal "google.golang.org/adk/internal/artifact" icontext "google.golang.org/adk/internal/context" "google.golang.org/adk/internal/toolinternal" "google.golang.org/adk/model" - "google.golang.org/adk/tool" "google.golang.org/adk/tool/loadartifactstool" ) @@ -277,7 +277,7 @@ func TestLoadArtifactsTool_ProcessRequest_Artifacts_OtherFunctionCall(t *testing } } -func createToolContext(t *testing.T) tool.Context { +func createToolContext(t *testing.T) agent.ToolContext { t.Helper() artifacts := &artifactinternal.Artifacts{ @@ -291,5 +291,5 @@ func createToolContext(t *testing.T) tool.Context { Artifacts: artifacts, }) - return toolinternal.NewToolContext(ctx, "", nil, nil) + return agent.NewToolContext(ctx, "", nil, nil) } diff --git a/tool/loadmemorytool/tool.go b/tool/loadmemorytool/tool.go index 6e29c143f..a0952e80c 100644 --- a/tool/loadmemorytool/tool.go +++ b/tool/loadmemorytool/tool.go @@ -22,12 +22,12 @@ import ( "google.golang.org/genai" + "google.golang.org/adk/agent" "google.golang.org/adk/internal/toolinternal" "google.golang.org/adk/internal/toolinternal/toolutils" "google.golang.org/adk/internal/utils" "google.golang.org/adk/memory" "google.golang.org/adk/model" - "google.golang.org/adk/tool" ) const memoryInstructions = `You have memory. You can use it to answer questions. If any questions need @@ -80,7 +80,7 @@ func (t *loadMemoryTool) Declaration() *genai.FunctionDeclaration { } // Run executes the tool with the provided context and arguments. -func (t *loadMemoryTool) Run(toolCtx tool.Context, args any) (map[string]any, error) { +func (t *loadMemoryTool) Run(toolCtx agent.ToolContext, args any) (map[string]any, error) { m, ok := args.(map[string]any) if !ok { return nil, fmt.Errorf("unexpected args type, got: %T", args) @@ -109,7 +109,7 @@ func (t *loadMemoryTool) Run(toolCtx tool.Context, args any) (map[string]any, er // ProcessRequest processes the LLM request by packing the tool and appending // memory-related instructions. -func (t *loadMemoryTool) ProcessRequest(ctx tool.Context, req *model.LLMRequest) error { +func (t *loadMemoryTool) ProcessRequest(ctx agent.ToolContext, req *model.LLMRequest) error { if err := toolutils.PackTool(req, t); err != nil { return err } diff --git a/tool/loadmemorytool/tool_test.go b/tool/loadmemorytool/tool_test.go index bee0082bb..5bde45722 100644 --- a/tool/loadmemorytool/tool_test.go +++ b/tool/loadmemorytool/tool_test.go @@ -22,12 +22,12 @@ import ( "google.golang.org/genai" + "google.golang.org/adk/agent" icontext "google.golang.org/adk/internal/context" "google.golang.org/adk/internal/toolinternal" "google.golang.org/adk/memory" "google.golang.org/adk/model" "google.golang.org/adk/session" - "google.golang.org/adk/tool" "google.golang.org/adk/tool/loadmemorytool" ) @@ -170,12 +170,12 @@ func TestLoadMemoryTool_ProcessRequest(t *testing.T) { } } -func createToolContext(t *testing.T, mem *mockMemory) tool.Context { +func createToolContext(t *testing.T, mem *mockMemory) agent.ToolContext { t.Helper() ctx := icontext.NewInvocationContext(t.Context(), icontext.InvocationContextParams{ Memory: mem, }) - return toolinternal.NewToolContext(ctx, "", nil, nil) + return agent.NewToolContext(ctx, "", nil, nil) } diff --git a/tool/mcptoolset/set_test.go b/tool/mcptoolset/set_test.go index 1fe3cb512..abe30593b 100644 --- a/tool/mcptoolset/set_test.go +++ b/tool/mcptoolset/set_test.go @@ -30,6 +30,7 @@ import ( "github.com/modelcontextprotocol/go-sdk/mcp" "google.golang.org/genai" + "google.golang.org/adk/agent" "google.golang.org/adk/agent/llmagent" icontext "google.golang.org/adk/internal/context" "google.golang.org/adk/internal/httprr" @@ -307,7 +308,7 @@ func TestCallToolReconnection(t *testing.T) { invCtx := icontext.NewInvocationContext(t.Context(), icontext.InvocationContextParams{}) ctx := icontext.NewReadonlyContext(invCtx) - toolCtx := toolinternal.NewToolContext(invCtx, "", nil, nil) + toolCtx := agent.NewToolContext(invCtx, "", nil, nil) // Get tools first to establish a session. tools, err := ts.Tools(ctx) @@ -415,6 +416,9 @@ func TestMCPToolSetConfirmation(t *testing.T) { city: "Lisbon", want: []*genai.Content{ genai.NewContentFromFunctionCall(toolName, map[string]any{"city": "Lisbon"}, "model"), + genai.NewContentFromFunctionResponse(toolName, map[string]any{ + "error": errors.New("error tool \"get_weather\" requires confirmation, please approve or reject"), + }, "user"), genai.NewContentFromFunctionCall(toolconfirmation.FunctionCallName, map[string]any{ "originalFunctionCall": &genai.FunctionCall{ Args: map[string]any{"city": "Lisbon"}, @@ -424,9 +428,6 @@ func TestMCPToolSetConfirmation(t *testing.T) { Hint: "Please approve or reject the tool call get_weather() by responding with a FunctionResponse with an expected ToolConfirmation payload.", }, }, "model"), - genai.NewContentFromFunctionResponse(toolName, map[string]any{ - "error": errors.New("error tool \"get_weather\" requires confirmation, please approve or reject"), - }, "user"), }, }, { @@ -438,6 +439,9 @@ func TestMCPToolSetConfirmation(t *testing.T) { confirmFunctionResponse: &genai.FunctionResponse{Name: toolconfirmation.FunctionCallName, Response: map[string]any{"confirmed": true}}, want: []*genai.Content{ genai.NewContentFromFunctionCall(toolName, map[string]any{"city": "Lisbon"}, "model"), + genai.NewContentFromFunctionResponse(toolName, map[string]any{ + "error": errors.New("error tool \"get_weather\" requires confirmation, please approve or reject"), + }, "user"), genai.NewContentFromFunctionCall(toolconfirmation.FunctionCallName, map[string]any{ "originalFunctionCall": &genai.FunctionCall{ Args: map[string]any{"city": "Lisbon"}, @@ -447,9 +451,6 @@ func TestMCPToolSetConfirmation(t *testing.T) { Hint: "Please approve or reject the tool call get_weather() by responding with a FunctionResponse with an expected ToolConfirmation payload.", }, }, "model"), - genai.NewContentFromFunctionResponse(toolName, map[string]any{ - "error": errors.New("error tool \"get_weather\" requires confirmation, please approve or reject"), - }, "user"), genai.NewContentFromFunctionResponse(toolName, map[string]any{ "output": map[string]any{"weather_summary": string(`Today in "Lisbon" is sunny`)}, }, "user"), @@ -465,6 +466,9 @@ func TestMCPToolSetConfirmation(t *testing.T) { confirmFunctionResponse: &genai.FunctionResponse{Name: toolconfirmation.FunctionCallName, Response: map[string]any{"confirmed": false}}, want: []*genai.Content{ genai.NewContentFromFunctionCall(toolName, map[string]any{"city": "Lisbon"}, "model"), + genai.NewContentFromFunctionResponse(toolName, map[string]any{ + "error": errors.New("error tool \"get_weather\" requires confirmation, please approve or reject"), + }, "user"), genai.NewContentFromFunctionCall(toolconfirmation.FunctionCallName, map[string]any{ "originalFunctionCall": &genai.FunctionCall{ Args: map[string]any{"city": "Lisbon"}, @@ -474,9 +478,6 @@ func TestMCPToolSetConfirmation(t *testing.T) { Hint: "Please approve or reject the tool call get_weather() by responding with a FunctionResponse with an expected ToolConfirmation payload.", }, }, "model"), - genai.NewContentFromFunctionResponse(toolName, map[string]any{ - "error": errors.New("error tool \"get_weather\" requires confirmation, please approve or reject"), - }, "user"), genai.NewContentFromFunctionResponse(toolName, map[string]any{ "error": errors.New("error tool \"get_weather\" call is rejected"), }, "user"), @@ -523,6 +524,9 @@ func TestMCPToolSetConfirmation(t *testing.T) { city: "Lisbon", want: []*genai.Content{ genai.NewContentFromFunctionCall(toolName, map[string]any{"city": "Lisbon"}, "model"), + genai.NewContentFromFunctionResponse(toolName, map[string]any{ + "error": errors.New("error tool \"get_weather\" requires confirmation, please approve or reject"), + }, "user"), genai.NewContentFromFunctionCall(toolconfirmation.FunctionCallName, map[string]any{ "originalFunctionCall": &genai.FunctionCall{ Args: map[string]any{"city": "Lisbon"}, @@ -532,9 +536,6 @@ func TestMCPToolSetConfirmation(t *testing.T) { Hint: "Please approve or reject the tool call get_weather() by responding with a FunctionResponse with an expected ToolConfirmation payload.", }, }, "model"), - genai.NewContentFromFunctionResponse(toolName, map[string]any{ - "error": errors.New("error tool \"get_weather\" requires confirmation, please approve or reject"), - }, "user"), }, }, { @@ -545,6 +546,9 @@ func TestMCPToolSetConfirmation(t *testing.T) { city: "Lisbon", want: []*genai.Content{ genai.NewContentFromFunctionCall(toolName, map[string]any{"city": "Lisbon"}, "model"), + genai.NewContentFromFunctionResponse(toolName, map[string]any{ + "error": errors.New("error tool \"get_weather\" requires confirmation, please approve or reject"), + }, "user"), genai.NewContentFromFunctionCall(toolconfirmation.FunctionCallName, map[string]any{ "originalFunctionCall": &genai.FunctionCall{ Args: map[string]any{"city": "Lisbon"}, @@ -554,9 +558,6 @@ func TestMCPToolSetConfirmation(t *testing.T) { Hint: "Please approve or reject the tool call get_weather() by responding with a FunctionResponse with an expected ToolConfirmation payload.", }, }, "model"), - genai.NewContentFromFunctionResponse(toolName, map[string]any{ - "error": errors.New("error tool \"get_weather\" requires confirmation, please approve or reject"), - }, "user"), }, }, { @@ -568,6 +569,9 @@ func TestMCPToolSetConfirmation(t *testing.T) { confirmFunctionResponse: &genai.FunctionResponse{Name: toolconfirmation.FunctionCallName, Response: map[string]any{"confirmed": true}}, want: []*genai.Content{ genai.NewContentFromFunctionCall(toolName, map[string]any{"city": "Lisbon"}, "model"), + genai.NewContentFromFunctionResponse(toolName, map[string]any{ + "error": errors.New("error tool \"get_weather\" requires confirmation, please approve or reject"), + }, "user"), genai.NewContentFromFunctionCall(toolconfirmation.FunctionCallName, map[string]any{ "originalFunctionCall": &genai.FunctionCall{ Args: map[string]any{"city": "Lisbon"}, @@ -577,9 +581,6 @@ func TestMCPToolSetConfirmation(t *testing.T) { Hint: "Please approve or reject the tool call get_weather() by responding with a FunctionResponse with an expected ToolConfirmation payload.", }, }, "model"), - genai.NewContentFromFunctionResponse(toolName, map[string]any{ - "error": errors.New("error tool \"get_weather\" requires confirmation, please approve or reject"), - }, "user"), genai.NewContentFromFunctionResponse(toolName, map[string]any{ "output": map[string]any{"weather_summary": string(`Today in "Lisbon" is sunny`)}, }, "user"), @@ -595,6 +596,9 @@ func TestMCPToolSetConfirmation(t *testing.T) { confirmFunctionResponse: &genai.FunctionResponse{Name: toolconfirmation.FunctionCallName, Response: map[string]any{"confirmed": false}}, want: []*genai.Content{ genai.NewContentFromFunctionCall(toolName, map[string]any{"city": "Lisbon"}, "model"), + genai.NewContentFromFunctionResponse(toolName, map[string]any{ + "error": errors.New("error tool \"get_weather\" requires confirmation, please approve or reject"), + }, "user"), genai.NewContentFromFunctionCall(toolconfirmation.FunctionCallName, map[string]any{ "originalFunctionCall": &genai.FunctionCall{ Args: map[string]any{"city": "Lisbon"}, @@ -604,9 +608,6 @@ func TestMCPToolSetConfirmation(t *testing.T) { Hint: "Please approve or reject the tool call get_weather() by responding with a FunctionResponse with an expected ToolConfirmation payload.", }, }, "model"), - genai.NewContentFromFunctionResponse(toolName, map[string]any{ - "error": errors.New("error tool \"get_weather\" requires confirmation, please approve or reject"), - }, "user"), genai.NewContentFromFunctionResponse(toolName, map[string]any{ "error": errors.New("error tool \"get_weather\" call is rejected"), }, "user"), diff --git a/tool/mcptoolset/tool.go b/tool/mcptoolset/tool.go index 7dbb45d5f..70b6d9f89 100644 --- a/tool/mcptoolset/tool.go +++ b/tool/mcptoolset/tool.go @@ -22,6 +22,7 @@ import ( "github.com/modelcontextprotocol/go-sdk/mcp" "google.golang.org/genai" + "google.golang.org/adk/agent" "google.golang.org/adk/internal/toolinternal" "google.golang.org/adk/internal/toolinternal/toolutils" "google.golang.org/adk/model" @@ -82,7 +83,7 @@ func (t *mcpTool) IsLongRunning() bool { return false } -func (t *mcpTool) ProcessRequest(ctx tool.Context, req *model.LLMRequest) error { +func (t *mcpTool) ProcessRequest(ctx agent.ToolContext, req *model.LLMRequest) error { return toolutils.PackTool(req, t) } @@ -90,7 +91,7 @@ func (t *mcpTool) Declaration() *genai.FunctionDeclaration { return t.funcDeclaration } -func (t *mcpTool) Run(ctx tool.Context, args any) (map[string]any, error) { +func (t *mcpTool) Run(ctx agent.ToolContext, args any) (map[string]any, error) { if confirmation := ctx.ToolConfirmation(); confirmation != nil { if !confirmation.Confirmed { return nil, fmt.Errorf("error tool %q %w", t.Name(), tool.ErrConfirmationRejected) diff --git a/tool/preloadmemorytool/tool.go b/tool/preloadmemorytool/tool.go index ab5866d85..16e4206c0 100644 --- a/tool/preloadmemorytool/tool.go +++ b/tool/preloadmemorytool/tool.go @@ -26,10 +26,10 @@ import ( "strings" "time" + "google.golang.org/adk/agent" "google.golang.org/adk/internal/utils" "google.golang.org/adk/memory" "google.golang.org/adk/model" - "google.golang.org/adk/tool" ) const preloadInstructions = `The following content is from your previous conversations with the user. @@ -71,7 +71,7 @@ func (t *preloadMemoryTool) IsLongRunning() bool { // ProcessRequest processes the LLM request by searching memory using the user's // current query and injecting relevant past conversations into system instructions. -func (t *preloadMemoryTool) ProcessRequest(ctx tool.Context, req *model.LLMRequest) error { +func (t *preloadMemoryTool) ProcessRequest(ctx agent.ToolContext, req *model.LLMRequest) error { userContent := ctx.UserContent() if userContent == nil || len(userContent.Parts) == 0 || userContent.Parts[0] == nil || userContent.Parts[0].Text == "" { diff --git a/tool/preloadmemorytool/tool_test.go b/tool/preloadmemorytool/tool_test.go index 20b77d98a..22fa3fd9a 100644 --- a/tool/preloadmemorytool/tool_test.go +++ b/tool/preloadmemorytool/tool_test.go @@ -23,12 +23,11 @@ import ( "google.golang.org/genai" + "google.golang.org/adk/agent" icontext "google.golang.org/adk/internal/context" - "google.golang.org/adk/internal/toolinternal" "google.golang.org/adk/memory" "google.golang.org/adk/model" "google.golang.org/adk/session" - "google.golang.org/adk/tool" "google.golang.org/adk/tool/preloadmemorytool" ) @@ -207,7 +206,7 @@ func TestPreloadMemoryTool_ProcessRequest(t *testing.T) { } } -func createToolContext(t *testing.T, mem *mockMemory, userContent *genai.Content) tool.Context { +func createToolContext(t *testing.T, mem *mockMemory, userContent *genai.Content) agent.ToolContext { t.Helper() ctx := icontext.NewInvocationContext(t.Context(), icontext.InvocationContextParams{ @@ -215,5 +214,5 @@ func createToolContext(t *testing.T, mem *mockMemory, userContent *genai.Content UserContent: userContent, }) - return toolinternal.NewToolContext(ctx, "", nil, nil) + return agent.NewToolContext(ctx, "", nil, nil) } diff --git a/tool/skilltoolset/internal/skilltool/list_skills.go b/tool/skilltoolset/internal/skilltool/list_skills.go index 88cae6923..1f24fb59c 100644 --- a/tool/skilltoolset/internal/skilltool/list_skills.go +++ b/tool/skilltoolset/internal/skilltool/list_skills.go @@ -18,6 +18,7 @@ import ( "html" "strings" + "google.golang.org/adk/agent" "google.golang.org/adk/tool" "google.golang.org/adk/tool/functiontool" "google.golang.org/adk/tool/skilltoolset/skill" @@ -38,13 +39,13 @@ func ListSkills(source skill.Source) (tool.Tool, error) { Name: "list_skills", Description: "Lists all available skills with their names and descriptions.", }, - func(ctx tool.Context, args ListSkillsArgs) (*ListSkillsResult, error) { + func(ctx agent.ToolContext, args ListSkillsArgs) (*ListSkillsResult, error) { return listSkills(ctx, args, source) }, ) } -func listSkills(ctx tool.Context, _ ListSkillsArgs, source skill.Source) (*ListSkillsResult, error) { +func listSkills(ctx agent.ToolContext, _ ListSkillsArgs, source skill.Source) (*ListSkillsResult, error) { skills, err := source.ListFrontmatters(ctx) if err != nil { return nil, err diff --git a/tool/skilltoolset/internal/skilltool/load_skill.go b/tool/skilltoolset/internal/skilltool/load_skill.go index 7e2a787e2..f662a8572 100644 --- a/tool/skilltoolset/internal/skilltool/load_skill.go +++ b/tool/skilltoolset/internal/skilltool/load_skill.go @@ -18,6 +18,7 @@ package skilltool import ( "fmt" + "google.golang.org/adk/agent" "google.golang.org/adk/tool" "google.golang.org/adk/tool/functiontool" "google.golang.org/adk/tool/skilltoolset/skill" @@ -51,13 +52,13 @@ func LoadSkill(source skill.Source) (tool.Tool, error) { Name: "load_skill", Description: "Loads the SKILL.md instructions for a given skill.", }, - func(ctx tool.Context, args LoadSkillArgs) (*LoadSkillResult, error) { + func(ctx agent.ToolContext, args LoadSkillArgs) (*LoadSkillResult, error) { return loadSkill(ctx, args, source) }, ) } -func loadSkill(ctx tool.Context, args LoadSkillArgs, source skill.Source) (*LoadSkillResult, error) { +func loadSkill(ctx agent.ToolContext, args LoadSkillArgs, source skill.Source) (*LoadSkillResult, error) { if args.Name == "" { return nil, fmt.Errorf("skill name is required to load a skill") } diff --git a/tool/skilltoolset/internal/skilltool/load_skill_resource.go b/tool/skilltoolset/internal/skilltool/load_skill_resource.go index 42abe0496..8eff8886c 100644 --- a/tool/skilltoolset/internal/skilltool/load_skill_resource.go +++ b/tool/skilltoolset/internal/skilltool/load_skill_resource.go @@ -18,6 +18,7 @@ import ( "fmt" "io" + "google.golang.org/adk/agent" "google.golang.org/adk/tool" "google.golang.org/adk/tool/functiontool" "google.golang.org/adk/tool/skilltoolset/skill" @@ -45,13 +46,13 @@ func LoadSkillResource(source skill.Source) (tool.Tool, error) { Name: "load_skill_resource", Description: "Loads a resource file (e.g., from references/ or assets/) associated with the specified skill.", }, - func(ctx tool.Context, args LoadSkillResourceArgs) (*LoadSkillResourceResult, error) { + func(ctx agent.ToolContext, args LoadSkillResourceArgs) (*LoadSkillResourceResult, error) { return loadSkillResource(ctx, args, source) }, ) } -func loadSkillResource(ctx tool.Context, args LoadSkillResourceArgs, source skill.Source) (*LoadSkillResourceResult, error) { +func loadSkillResource(ctx agent.ToolContext, args LoadSkillResourceArgs, source skill.Source) (*LoadSkillResourceResult, error) { if args.SkillName == "" { return nil, fmt.Errorf("skill name is required to load a resource") } diff --git a/tool/skilltoolset/internal/skilltool/tools_test.go b/tool/skilltoolset/internal/skilltool/tools_test.go index d181e55f3..3c1f37275 100644 --- a/tool/skilltoolset/internal/skilltool/tools_test.go +++ b/tool/skilltoolset/internal/skilltool/tools_test.go @@ -22,9 +22,9 @@ import ( "github.com/google/go-cmp/cmp" + "google.golang.org/adk/agent" icontext "google.golang.org/adk/internal/context" "google.golang.org/adk/internal/toolinternal" - "google.golang.org/adk/tool" "google.golang.org/adk/tool/skilltoolset/internal/skilltool" "google.golang.org/adk/tool/skilltoolset/skill" ) @@ -74,9 +74,9 @@ func (m *mockSource) LoadResource(ctx context.Context, name, resourcePath string return io.NopCloser(strings.NewReader(res)), nil } -func createToolContext(t *testing.T) tool.Context { +func createToolContext(t *testing.T) agent.ToolContext { invCtx := icontext.NewInvocationContext(t.Context(), icontext.InvocationContextParams{}) - return toolinternal.NewToolContext(invCtx, "", nil, nil) + return agent.NewToolContext(invCtx, "", nil, nil) } func TestListSkills(t *testing.T) { diff --git a/tool/skilltoolset/toolset.go b/tool/skilltoolset/toolset.go index 3e681ad16..7754a9fbd 100644 --- a/tool/skilltoolset/toolset.go +++ b/tool/skilltoolset/toolset.go @@ -104,7 +104,7 @@ func (ts *SkillToolset) Tools(ctx agent.ReadonlyContext) ([]tool.Tool, error) { // ProcessRequest implements toolinternal.RequestProcessor. It attaches // the list of available skills and the system instruction explaining to the // agent what it can do with these skills. -func (ts *SkillToolset) ProcessRequest(ctx tool.Context, req *model.LLMRequest) error { +func (ts *SkillToolset) ProcessRequest(ctx agent.ToolContext, req *model.LLMRequest) error { skills, err := ts.source.ListFrontmatters(ctx) if err != nil { return err diff --git a/tool/tool.go b/tool/tool.go index 19a72ac8e..091df8883 100644 --- a/tool/tool.go +++ b/tool/tool.go @@ -18,7 +18,6 @@ package tool import ( - "context" "errors" "fmt" @@ -26,10 +25,7 @@ import ( "google.golang.org/adk/agent" "google.golang.org/adk/internal/toolinternal/toolutils" - "google.golang.org/adk/memory" "google.golang.org/adk/model" - "google.golang.org/adk/session" - "google.golang.org/adk/tool/toolconfirmation" ) // ErrConfirmationRequired indicates that the tool requires confirmation. @@ -49,57 +45,12 @@ type Tool interface { IsLongRunning() bool } -// Context defines the interface for the context passed to a tool when it's -// called. It provides access to invocation-specific information and allows -// the tool to interact with the agent's state and memory. -type Context interface { - agent.CallbackContext - // FunctionCallID returns the unique identifier of the function call - // that triggered this tool execution. - FunctionCallID() string - - // Actions returns the EventActions for the current event. This can be - // used by the tool to modify the agent's state, transfer to another - // agent, or perform other actions. - Actions() *session.EventActions - // SearchMemory performs a semantic search on the agent's memory. - SearchMemory(context.Context, string) (*memory.SearchResponse, error) - - // ToolConfirmation returns a handler for checking the Human-in-the-Loop - // confirmation status for the current tool context. This should be used within a tool's logic - // *before* performing any sensitive operations that require user approval. - // - // Example Usage: - // if confirmation := ctx.ToolConfirmation(); confirmation == nil { - // // Confirmation required, create confirmation or handle appropriately - // ctx.RequestConfirmation("hint", payload) - // } - // - // The returned *toolconfirmation.ToolConfirmation object provides methods to check the actual - // confirmation state. - ToolConfirmation() *toolconfirmation.ToolConfirmation - - // RequestConfirmation initiates the Human-in-the-Loop (HITL) process to ask the user for approval - // before the tool proceeds with a specific action. Call this method when a tool needs - // explicit user consent. - // - // This will typically result in the ADK emitting a special event - // (e.g., a FunctionCall like "adk_request_confirmation") to the client application/UI, - // prompting the user for a decision. - // - // Args: - // - hint: A human-readable string explaining why confirmation is needed. This is usually - // displayed to the user in the confirmation prompt. - // - payload: Any additional data or context about the action requiring confirmation. - // - // Returns: - // - nil: If the confirmation request was successfully enqueued or initiated within the ADK. - // This indicates that the process of asking the user has begun. It does NOT mean the action - // is approved. The tool's execution will likely pause or be suspended until the user responds. - // - error: If there was a failure in initiating the confirmation process itself (e.g., invalid - // arguments, issue with the event system). The request to ask the user has not been sent. - RequestConfirmation(hint string, payload any) error -} +// Context is an alias for agent.ToolContext. +// +// Deprecated: use agent.Context directly. This alias exists only to +// minimize churn during the migration and will be removed in a future +// release. +type Context = agent.ToolContext // Toolset is an interface for a collection of tools. It allows grouping // related tools together and providing them to an agent. @@ -238,18 +189,18 @@ type confirmationTool struct { type runnableTool interface { Tool Declaration() *genai.FunctionDeclaration - Run(ctx Context, args any) (result map[string]any, err error) + Run(ctx agent.ToolContext, args any) (result map[string]any, err error) } func (t *confirmationTool) Declaration() *genai.FunctionDeclaration { return t.runnableTool.Declaration() } -func (t *confirmationTool) ProcessRequest(ctx Context, req *model.LLMRequest) error { +func (t *confirmationTool) ProcessRequest(ctx agent.ToolContext, req *model.LLMRequest) error { return toolutils.PackTool(req, t) } -func (t *confirmationTool) Run(ctx Context, args any) (map[string]any, error) { +func (t *confirmationTool) Run(ctx agent.ToolContext, args any) (map[string]any, error) { ft := t.runnableTool // Check for Human-in-the-Loop confirmation. diff --git a/tool/tool_test.go b/tool/tool_test.go index 3f5aa870e..4bc9ccce1 100644 --- a/tool/tool_test.go +++ b/tool/tool_test.go @@ -54,7 +54,7 @@ func TestTypes(t *testing.T) { { name: "FunctionTool", constructor: func() (tool.Tool, error) { - return functiontool.New(functiontool.Config{}, func(_ tool.Context, input intInput) (intOutput, error) { + return functiontool.New(functiontool.Config{}, func(_ agent.ToolContext, input intInput) (intOutput, error) { return intOutput(input), nil }) }, @@ -160,7 +160,7 @@ func (tts *testToolset) Tools(agent.ReadonlyContext) ([]tool.Tool, error) { func TestWithConfirmation(t *testing.T) { toolRan := false - noOpTool, err := functiontool.New(functiontool.Config{Name: "noOpTool"}, func(ctx tool.Context, input struct{}) (struct{}, error) { + noOpTool, err := functiontool.New(functiontool.Config{Name: "noOpTool"}, func(ctx agent.ToolContext, input struct{}) (struct{}, error) { toolRan = true return struct{}{}, nil }) diff --git a/util/aiplatform/aiplatform.go b/util/aiplatform/aiplatform.go new file mode 100644 index 000000000..dfb753c9a --- /dev/null +++ b/util/aiplatform/aiplatform.go @@ -0,0 +1,31 @@ +// Copyright 2026 Google LLC +// +// Licensed under the Apache License, Version 2.0 (the "License"); +// you may not use this file except in compliance with the License. +// You may obtain a copy of the License at +// +// http://www.apache.org/licenses/LICENSE-2.0 +// +// Unless required by applicable law or agreed to in writing, software +// distributed under the License is distributed on an "AS IS" BASIS, +// WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. +// See the License for the specific language governing permissions and +// limitations under the License. + +// Package aiplatform provides utils for aiplatform.googleapis.com +package aiplatform + +// HostURL returns aiplatform host url for a given region +// Example: "us-central1-aiplatform.googleapis.com" +func HostURL(region string) string { + if region == "" || region == "global" { + return "aiplatform.googleapis.com" + } + return region + "-aiplatform.googleapis.com" +} + +// HostPortUrl returns aiplatform host url for a given region with port number +// Example: "us-central1-aiplatform.googleapis.com:443" +func HostPortURL(region string) string { + return HostURL(region) + ":443" +} diff --git a/util/vertexai/agent_engine.go b/util/vertexai/agent_engine.go new file mode 100644 index 000000000..69bf10f7e --- /dev/null +++ b/util/vertexai/agent_engine.go @@ -0,0 +1,41 @@ +// Copyright 2026 Google LLC +// +// Licensed under the Apache License, Version 2.0 (the "License"); +// you may not use this file except in compliance with the License. +// You may obtain a copy of the License at +// +// http://www.apache.org/licenses/LICENSE-2.0 +// +// Unless required by applicable law or agreed to in writing, software +// distributed under the License is distributed on an "AS IS" BASIS, +// WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. +// See the License for the specific language governing permissions and +// limitations under the License. + +// Package vertexai provides utilities for Agent Engine deployments +package vertexai + +import "fmt" + +const ( + agentEngineTemplate = "projects/%s/locations/%s/reasoningEngines/%s" + agentEngineSessionTemplate = agentEngineTemplate + "/sessions/%s" +) + +type AgentEngineData struct { + Location string + ProjectID string + ReasoningEngine string +} + +// AgentEngineResource returns a formatted string indicating agent engine instance +// (template `projects/%s/locations/%s/reasoningEngines/%s`) +func AgentEngineResource(data *AgentEngineData) string { + return fmt.Sprintf(agentEngineTemplate, data.ProjectID, data.Location, data.ReasoningEngine) +} + +// SessionResource returns a formatted string indicating specific session for an agent engine instance +// (template `projects/%s/locations/%s/reasoningEngines/%s/sessions/%s`) +func SessionResource(data *AgentEngineData, sessionID string) string { + return fmt.Sprintf(agentEngineSessionTemplate, data.ProjectID, data.Location, data.ReasoningEngine, sessionID) +}