diff --git a/README.md b/README.md index 65d51ff..a2ebbf3 100644 --- a/README.md +++ b/README.md @@ -9,6 +9,9 @@ Responses agent as an To run the supervisor and ACP child inside Foundry itself, use the [Hosted Agent v2 package and Kubernetes gateway](docs/foundry-hosted-v2.md). +For AgentKit agents that use Orka tools, configure the +[AgentKit continuation secret](docs/harness-v2.md#agentkit-tool-workflows) on the +broker and hosted agent. The adapter calls the Hosted Agent's dedicated Responses endpoint: diff --git a/docs/harness-v2.md b/docs/harness-v2.md index 45aa83b..52db512 100644 --- a/docs/harness-v2.md +++ b/docs/harness-v2.md @@ -83,6 +83,7 @@ configuration digest, plus: | `ORKA_FOUNDRY_BROKER_ADDR` | `127.0.0.1:8091`. | | `ORKA_FOUNDRY_BROKER_STATE_DIR` | Absolute path on the broker-only persistent volume. | | `ORKA_FOUNDRY_BROKER_BEARER_TOKEN` | At least 32 bytes, from a Kubernetes Secret. | +| `ORKA_FOUNDRY_BROKER_AGENTKIT_CONTINUATION_PROOF` | Optional shared secret for an AgentKit Hosted Agent's governed tool results. See below. | Only the broker receives Azure Workload Identity or another refreshable `DefaultAzureCredential` configuration. The initial implementation requires @@ -106,6 +107,53 @@ contains remote identifiers and ownership metadata, never prompts, tool arguments, provider response bodies, Azure tokens, or local bearer tokens. Do not delete or replace this ledger while it owns remote work. +## AgentKit tool workflows + +Configure the hosted AgentKit agent with static `brokeredTools` and +`AGENTKIT_FOUNDRY_BROKERED_MODEL_LOOP=1`. Use an AgentKit version that supports +sequential tool rounds and set `toolSchemaMode` to `provider-static` in the +Foundry configuration. The hosted agent and Orka runtime must use the same +tool names. Each round proposes one tool; Orka executes it and the hosted agent +receives its result before deciding whether to call another tool or answer. + +For an AgentKit Hosted Agent configured with brokered tools, give the broker +`ORKA_FOUNDRY_BROKER_AGENTKIT_CONTINUATION_PROOF` and give the hosted AgentKit +process `AGENTKIT_FOUNDRY_BROKERED_CONTINUATION_PROOF` with the same value. Use a +secret of at least 32 bytes, without whitespace or control characters. Keep it +in secret-backed environment variables for those two processes only. Never +put it in `foundry.json`, an image, an Orka Task, the supervisor environment, +or the ACP child environment. Use a separate secret for each deployment pair. + +When the agent requests a tool, the ACP child calls Orka's session MCP server. +The broker validates the resulting request against its owner, prompt, lease, +previous response and pending call IDs before returning the result to AgentKit. +It then attaches the secret in the top-level `brokered_continuation_proof` JSON +field. Ordinary prompts do not carry it. The broker rejects proof fields or +proof headers supplied by its caller and never persists the secret in its +ledger or returns it to the ACP child. + +AgentKit expects a result envelope with `approved` and either `output` or +`error`. The broker converts Orka's validated MCP result to this format. +An explicit `isError: false` becomes `approved: true` with the text and +structured content preserved under `output`. `isError: true` becomes +`approved: false` with a `brokered_tool_error` code and the validated error +text. Malformed results and missing error flags are rejected. MCP authorization +failures abort the tool call before any result is sent to AgentKit. The +`approved` field is AgentKit's result format; it does not report human approval. +Approval-required tools remain unsupported by this adapter. + +This is the existing AgentKit shared-secret contract. It authenticates the +broker's continuation route; it is not a signed execution receipt. The broker's +ownership checks and AgentKit's session, pending-call and replay checks remain +required. The Foundry Responses gateway must preserve the body extension and +deliver it to the hosted wrapper. Local contract tests do not establish that +the deployed Foundry gateway forwards it; verify the configured agent version +before enabling its tools in a live runtime. + +Leave the variable unset for other Hosted Agents. Their function output format +and requests remain unchanged. AgentKit must also support repeated tool rounds +for workflows that need several lookups before answering. + ## Lifecycle guarantees and limits The broker persists a random caller-chosen remote session ID before sending @@ -135,6 +183,16 @@ remote execution can survive the local container. Preserve unresolved ownership records for investigation; do not fabricate retirement receipts or remove finalizers to bypass them. +The broker writes one bounded JSON diagnostic to stderr when a dispatched +response fails. It records the failure stage, outer HTTP status, invocation +sequence, hashed owner, and whether a response acknowledgement was persisted. +An observed terminal frame adds its status and an allowlisted AgentKit error +code. The optional `error.upstream_status` is recorded only as an integer from +400 through 599. Unknown error codes become `unknown`; provider messages, +response bodies, URLs, remote IDs, headers, and credentials are excluded. +These diagnostics leave ownership and cleanup decisions to the existing +durable evidence and settlement checks. + An authenticated drain can replace a surviving supervisor after a controller epoch change. After a supervisor crash, Orka cannot import the broker's old-owner proof through the current harness contract. That recovery remains @@ -148,3 +206,41 @@ tool allowlists, and blocked output. Broker tests cover durable ownership, lease cleanup, repeated controls, remote stop/delete proof, and ambiguous acceptance. Live validation additionally requires the exact configured Hosted Agent version and Azure identity. + +For local tests against an AgentKit source checkout, install its common package +in a Python environment and run the opt-in integration tests: + +```sh +export AGENTKIT_SOURCE_DIR=/path/to/agentkit +uv venv /tmp/foundry-agentkit-venv +export AGENTKIT_PYTHON=/tmp/foundry-agentkit-venv/bin/python +uv pip install --python "$AGENTKIT_PYTHON" -e "$AGENTKIT_SOURCE_DIR/runtimes/common" +go test ./internal/broker -run TestBrokerAgentKitHosted -count=1 -v +``` + +This runs the production hosted AgentKit server and model loop, the Foundry ACP +entrypoint over pipes, and the lifecycle broker. It checks two sequential tools, +tool-error recovery, authorization denial, response identity changes, and proof +isolation. A gateway that strips the proof is also tested to verify that no model +resume occurs. Cancellation cases hold the model connection open, wait for an +early hosted response acknowledgement, and verify that disconnect and lease +expiry close the model connection and allow proven retirement. A gateway that +loses the acknowledgement must leave the broker's ownership unresolved. + +To include the native Microsoft Agent Framework path without brokered tools, +install its adapter and select that Python environment as well: + +```sh +uv pip install --python "$AGENTKIT_PYTHON" -e "$AGENTKIT_SOURCE_DIR/runtimes/microsoft-agent-framework" +export AGENTKIT_MAF_PYTHON="$AGENTKIT_PYTHON" +go test ./internal/broker -run TestBrokerAgentKitHostedCancellation/native_disconnect -count=1 -v +``` + +This case uses the real MAF runtime against a held model connection and verifies +that cancellation closes that connection after the broker records the early +response ID. It is skipped when `AGENTKIT_MAF_PYTHON` is unset. + +The model, MCP backend, supervisor context stamping, and Azure +session-management API are local fixtures. It requires no Azure or model credentials +and does not validate a deployed Orka controller or the public Foundry gateway. +These tests are skipped when `AGENTKIT_SOURCE_DIR` is unset. diff --git a/internal/acp/mcp.go b/internal/acp/mcp.go index 2f15897..d2e5141 100644 --- a/internal/acp/mcp.go +++ b/internal/acp/mcp.go @@ -123,7 +123,7 @@ func (m *acpMCPClient) execute(ctx context.Context, name string, args json.RawMe IsError *bool `json:"isError,omitempty"` StructuredContent json.RawMessage `json:"structuredContent,omitempty"` } - if strictjson.Decode(result, &reply, false) != nil || reply.Content == nil { + if strictjson.DecodeStruct(result, &reply, false) != nil || reply.Content == nil { return "", false, errACPMCP } for _, content := range reply.Content { diff --git a/internal/acp/mcp_metadata_test.go b/internal/acp/mcp_metadata_test.go index dfa28fc..7ee591d 100644 --- a/internal/acp/mcp_metadata_test.go +++ b/internal/acp/mcp_metadata_test.go @@ -2,6 +2,7 @@ package acp import ( "encoding/json" + "fmt" "net/http" "reflect" "strings" @@ -11,6 +12,27 @@ import ( "github.com/orka-agents/agent-runtime-foundry/internal/foundry" ) +func TestACPToolOutputRejectsConflictingFoldedErrorFlags(t *testing.T) { + for _, flags := range []string{`"isError":true,"ISERROR":false`, `"ISERROR":true,"isError":false`} { + t.Run(flags, func(t *testing.T) { + mcp := &acpTestMCP{tools: func() []map[string]any { return acpTestTools("probe") }} + mcp.execute = func(w http.ResponseWriter, _ *http.Request, id json.RawMessage, _ string, _ json.RawMessage) { + w.Header().Set("Content-Type", "application/json") + _, _ = fmt.Fprintf(w, `{"jsonrpc":"2.0","id":%s,"result":{"content":[{"type":"text","text":"tool failed"}],%s}}`, id, flags) + } + var requests atomic.Int32 + peer := newACPTestPeer(t, foundry.ToolSchemaModeProviderStatic, func(w http.ResponseWriter, _ *http.Request) { + requests.Add(1) + acpTestCompleted(w, "tool-response", "", acpTestCall("probe", "call-probe", `{}`)) + }, mcp) + acpAssertFailure(t, peer.reply(peer.prompt("use probe"))) + if requests.Load() != 1 || mcp.calls.Load() != 1 || acpOutput(peer.events) != "" { + t.Fatal("contradictory MCP error flags were converted into a successful continuation") + } + }) + } +} + func TestACPToolOutputForwardsOnlyValidatedModelContent(t *testing.T) { for _, mode := range []string{"text", "structured", "error"} { t.Run(mode, func(t *testing.T) { diff --git a/internal/broker/agentkit.go b/internal/broker/agentkit.go new file mode 100644 index 0000000..e1550db --- /dev/null +++ b/internal/broker/agentkit.go @@ -0,0 +1,97 @@ +package broker + +import ( + "encoding/json" + "strings" + "unicode" + "unicode/utf8" + + "github.com/orka-agents/agent-runtime-foundry/internal/foundry" + "github.com/orka-agents/agent-runtime-foundry/internal/strictjson" +) + +const ( + brokerAgentKitProofEnv = "ORKA_FOUNDRY_BROKER_AGENTKIT_CONTINUATION_PROOF" + brokerAgentKitProofHeader = "X-AgentKit-Brokered-Continuation-Proof" +) + +func brokerAgentKitProofValid(value string) bool { + return value == "" || (len(value) >= 32 && utf8.ValidString(value) && foundry.SafeString(value, 16<<10) && strings.IndexFunc(value, unicode.IsSpace) < 0) +} + +// The proof never enters the shared request type used by the ACP child. This +// wire-only extension is added after the broker's owner, lease and call checks. +func (b *lifecycleBroker) marshalResponseRequest(request foundry.ResponseRequest) ([]byte, error) { + proof := "" + if inputs, ok := request.Input.([]any); ok && len(inputs) != 0 { + proof = b.cfg.agentKitProof + } + return json.Marshal(struct { + foundry.ResponseRequest + ContinuationProof string `json:"brokered_continuation_proof,omitempty"` + }{ResponseRequest: request, ContinuationProof: proof}) +} + +// Orka's MCP proxy explicitly reports isError, with text content and optional +// structured output. Reject other envelopes rather than infer authorization +// from arbitrary tool data. JSON-RPC authorization failures never reach here. +func brokerAgentKitOutputs(request *foundry.ResponseRequest) error { + if _, first := request.Input.(string); first { + return nil + } + inputs, ok := request.Input.([]any) + if !ok || len(inputs) == 0 { + return errBrokerInvalid + } + for _, input := range inputs { + item, ok := input.(map[string]any) + if !ok { + return errBrokerInvalid + } + output, ok := item["output"].(string) + if !ok { + return errBrokerInvalid + } + var result struct { + Content []struct { + Type string `json:"type"` + Text *string `json:"text"` + } `json:"content"` + IsError *bool `json:"isError"` + StructuredContent json.RawMessage `json:"structuredContent,omitempty"` + } + if strictjson.DecodeStruct([]byte(output), &result, true) != nil || result.Content == nil || result.IsError == nil { + return errBrokerInvalid + } + var message strings.Builder + for i, content := range result.Content { + if content.Type != "text" || content.Text == nil { + return errBrokerInvalid + } + if i != 0 { + message.WriteByte('\n') + } + message.WriteString(*content.Text) + } + if len(result.StructuredContent) != 0 && result.StructuredContent[0] != '{' { + return errBrokerInvalid + } + var normalized any + if *result.IsError { + if message.Len() == 0 { + message.WriteString("The governed tool returned an error.") + } + normalized = map[string]any{"approved": false, "error": map[string]string{ + "code": "brokered_tool_error", "message": message.String(), + }} + } else { + normalized = map[string]any{"approved": true, "output": result} + } + encoded, err := json.Marshal(normalized) + if err != nil || len(encoded) > foundry.DefaultMaxBrokeredBytes { + return errBrokerInvalid + } + item["output"] = string(encoded) + } + return nil +} diff --git a/internal/broker/agentkit_cancellation_test.go b/internal/broker/agentkit_cancellation_test.go new file mode 100644 index 0000000..e8de7b0 --- /dev/null +++ b/internal/broker/agentkit_cancellation_test.go @@ -0,0 +1,176 @@ +package broker + +import ( + "bufio" + "bytes" + "context" + "io" + "net/http" + "net/http/httptest" + "net/url" + "os" + "path/filepath" + "strings" + "sync/atomic" + "testing" + "time" + + "github.com/orka-agents/agent-runtime-foundry/internal/brokerapi" +) + +// Exercise cancellation while the real hosted AgentKit is awaiting a model. +// Azure's gateway and session lifecycle remain fixtures; a lost acknowledgement +// must stay unresolved even when the fixture reports that compute is idle. +func TestBrokerAgentKitHostedCancellation(t *testing.T) { + source := os.Getenv("AGENTKIT_SOURCE_DIR") + if source == "" { + t.Skip("set AGENTKIT_SOURCE_DIR and AGENTKIT_PYTHON to test an AgentKit checkout") + } + for _, mode := range []string{"disconnect", "lease_expiry", "acknowledgement_lost", "native_disconnect"} { + t.Run(mode, func(t *testing.T) { + native := mode == "native_disconnect" + if native && os.Getenv("AGENTKIT_MAF_PYTHON") == "" { + t.Skip("set AGENTKIT_MAF_PYTHON to test the native Microsoft Agent Framework runtime") + } + modelStarted, modelCancelled := make(chan struct{}), make(chan struct{}) + var models, inferences atomic.Int32 + model := httptest.NewServer(http.HandlerFunc(func(w http.ResponseWriter, r *http.Request) { + body, err := io.ReadAll(r.Body) + if err != nil || r.Method != http.MethodPost || r.URL.Path != "/v1/chat/completions" || + bytes.Contains(body, []byte(brokerAgentKitFixtureProof)) || r.Header.Get(brokerAgentKitProofHeader) != "" { + t.Error("model request used the wrong route or exposed continuation credentials") + w.WriteHeader(http.StatusBadRequest) + return + } + if models.Add(1) != 1 { + t.Error("cancelled inference was replayed") + w.WriteHeader(http.StatusConflict) + return + } + close(modelStarted) + // Never finish the model response. Socket closure, rather than a + // mocked CancelledError, must interrupt AgentKit's model request. + <-r.Context().Done() + close(modelCancelled) + })) + t.Cleanup(model.Close) + hostedURL, stateFile := brokerStartAgentKit(t, source, model.URL, native) + f := newBrokerFixture(t, "success") + cfg := brokerTestConfig(t, f) + cfg.agentKitProof = brokerAgentKitFixtureProof + acknowledgementLost := make(chan struct{}) + client := &http.Client{Transport: brokerFixtureTransport(func(r *http.Request) (*http.Response, error) { + if !strings.HasSuffix(r.URL.Path, "/endpoint/protocols/openai/responses") { + return http.DefaultTransport.RoundTrip(r) + } + if inferences.Add(1) != 1 { + t.Error("cancelled hosted request was resubmitted") + return nil, errBrokerConflict + } + request := r.Clone(r.Context()) + request.URL, _ = url.Parse(hostedURL + "/responses") + request.Host = request.URL.Host + response, err := http.DefaultTransport.RoundTrip(request) + if err != nil || mode != "acknowledgement_lost" { + return response, err + } + // Simulate a gateway losing the first frame. An acknowledgement + // accepted only by the gateway does not prove broker ownership. + reader := bufio.NewReader(response.Body) + var frame bytes.Buffer + for { + line, readErr := reader.ReadString('\n') + if readErr != nil || frame.Len()+len(line) > 1<<20 { + _ = response.Body.Close() + return nil, errBrokerAmbiguous + } + frame.WriteString(line) + if strings.TrimSpace(line) == "" { + break + } + } + if !strings.HasPrefix(response.Header.Get("Content-Type"), "text/event-stream") || + !bytes.Contains(frame.Bytes(), []byte(`"response.created"`)) { + t.Error("hosted server did not send an early response acknowledgement") + _ = response.Body.Close() + return nil, errBrokerAmbiguous + } + response.Body = struct { + io.Reader + io.Closer + }{reader, response.Body} + close(acknowledgementLost) + return response, nil + })} + b, server := startBrokerTestWithClient(t, cfg, client) + c := brokerTestContext(cfg) + if mode == "lease_expiry" { + c.LeaseExpiresAt = time.Now().Add(2 * time.Second).UTC().Format(time.RFC3339Nano) + } + ctx, cancel := context.WithCancel(context.Background()) + t.Cleanup(cancel) + done := brokerAsyncInference(ctx, server.URL, c, brokerTestBody("")) + select { + case <-modelStarted: + case <-time.After(4 * time.Second): + t.Fatal("hosted model request did not start") + } + if mode == "acknowledgement_lost" { + select { + case <-acknowledgementLost: + case <-time.After(4 * time.Second): + t.Fatal("gateway did not receive the early acknowledgement") + } + if brokerInvocationState(b, c) != "intent" { + t.Fatal("gateway-only acknowledgement was admitted as broker evidence") + } + } else { + brokerAwait(t, func() bool { return brokerInvocationState(b, c) == "accepted" }) + } + if mode != "lease_expiry" { + cancel() + } + if result := brokerWaitInference(t, done); result.err == nil && result.status == http.StatusOK { + t.Fatal("cancelled hosted inference exposed a successful terminal response") + } + if mode == "acknowledgement_lost" { + brokerAwait(t, func() bool { return brokerInvocationState(b, c) == "uncertain" }) + brokerPendingControl(t, server.URL, brokerapi.SettlePath, c, false, 1) + brokerPendingControl(t, server.URL, brokerapi.RetirePath, c, false, 1) + } else { + proof := brokerTestControl(t, server.URL, brokerapi.SettlePath, c) + if !proof.SettlementProven || proof.ActiveInvocations != 0 || proof.AmbiguousInvocations != 0 { + t.Fatal("acknowledged hosted cancellation did not settle") + } + if proof := brokerTestControl(t, server.URL, brokerapi.RetirePath, c); !proof.RetirementProven { + t.Fatal("acknowledged hosted owner did not retire") + } + } + select { + case <-modelCancelled: + case <-time.After(4 * time.Second): + t.Fatal("hosted client disconnect left the model request running") + } + brokerAwait(t, func() bool { _, _, stops, _ := f.counts(); return stops > 0 }) + creates, _, _, deletes := f.counts() + wantDeletes := 1 + if mode == "acknowledgement_lost" { + wantDeletes = 0 + } + if creates != 1 || inferences.Load() != 1 || models.Load() != 1 || deletes != wantDeletes { + t.Fatal("hosted cancellation replayed work or deleted an unresolved owner") + } + for _, path := range []string{stateFile, filepath.Join(cfg.stateDir, "state.json")} { + data, err := os.ReadFile(path) + // An initial cancelled request has no hosted continuation to + // cache. The broker must always retain its ownership ledger. + if path == stateFile && os.IsNotExist(err) { + continue + } + if err != nil || bytes.Contains(data, []byte(cfg.agentKitProof)) { + t.Fatal("durable state was missing or retained continuation credentials") + } + } + }) + } +} diff --git a/internal/broker/agentkit_integration_test.go b/internal/broker/agentkit_integration_test.go new file mode 100644 index 0000000..30feecb --- /dev/null +++ b/internal/broker/agentkit_integration_test.go @@ -0,0 +1,426 @@ +package broker + +import ( + "bytes" + "context" + "encoding/base64" + "encoding/json" + "fmt" + "io" + "net" + "net/http" + "net/http/httptest" + "net/url" + "os" + "os/exec" + "path/filepath" + "strings" + "sync" + "sync/atomic" + "testing" + "time" + + "github.com/orka-agents/agent-runtime-foundry/internal/acp" + "github.com/orka-agents/agent-runtime-foundry/internal/brokerapi" + "github.com/orka-agents/agent-runtime-foundry/internal/foundry" +) + +// This opt-in test runs AgentKit's production hosted server and model loop, +// Foundry's ACP entrypoint over real pipes, and the real lifecycle broker. +// Only the model, MCP backend, supervisor context stamping, and Azure session +// management API are fixtures. It does not establish Azure ingress support. +func TestBrokerAgentKitHostedIntegration(t *testing.T) { + source := os.Getenv("AGENTKIT_SOURCE_DIR") + if source == "" { + t.Skip("set AGENTKIT_SOURCE_DIR and AGENTKIT_PYTHON to test an AgentKit checkout") + } + for _, mode := range []string{"success", "tool_error", "authorization_denied", "gateway_strips_proof"} { + t.Run(mode, func(t *testing.T) { + var models, calls, invocations atomic.Int32 + model := httptest.NewServer(http.HandlerFunc(func(w http.ResponseWriter, r *http.Request) { + body, _ := io.ReadAll(r.Body) + if r.Method != http.MethodPost || r.URL.Path != "/v1/chat/completions" || + bytes.Contains(body, []byte(brokerAgentKitFixtureProof)) || r.Header.Get(brokerAgentKitProofHeader) != "" { + t.Error("model request used the wrong route or exposed continuation credentials") + w.WriteHeader(http.StatusBadRequest) + return + } + var request struct { + Messages []map[string]any `json:"messages"` + Tools []map[string]any `json:"tools"` + } + if json.Unmarshal(body, &request) != nil || len(request.Tools) != 2 { + t.Error("hosted model loop lost its static tool schemas") + w.WriteHeader(http.StatusBadRequest) + return + } + n := models.Add(1) + if n > 1 { + last := request.Messages[len(request.Messages)-1] + text, _ := last["content"].(string) + var result struct { + Approved *bool `json:"approved"` + Output json.RawMessage `json:"output"` + Error map[string]any `json:"error"` + } + approved := mode != "tool_error" || n != 2 + if last["role"] != "tool" || json.Unmarshal([]byte(text), &result) != nil || + result.Approved == nil || *result.Approved != approved || + (!approved && (result.Output != nil || result.Error["code"] != "brokered_tool_error")) { + t.Error("real model resume lost the governed tool result or converted an error to approval") + w.WriteHeader(http.StatusBadRequest) + return + } + } + message := map[string]any{"role": "assistant", "content": "Checked the work order and inventory."} + if n < 3 { + name, arguments := "lookup_work_order", `{"id":"WO-7"}` + if n == 2 { + name, arguments = "lookup_inventory", `{"part":"FBR-7"}` + } + message = map[string]any{"role": "assistant", "content": nil, "tool_calls": []any{map[string]any{ + "id": "model-chosen-id", "type": "function", "function": map[string]string{"name": name, "arguments": arguments}, + }}} + } + w.Header().Set("Content-Type", "application/json") + _ = json.NewEncoder(w).Encode(map[string]any{"choices": []any{map[string]any{"message": message}}}) + })) + t.Cleanup(model.Close) + hostedURL, stateFile := brokerStartAgentKit(t, source, model.URL, false) + f := newBrokerFixture(t, "success") + cfg := brokerTestConfig(t, f) + cfg.agentKitProof = brokerAgentKitFixtureProof + var remoteMu sync.Mutex + var remoteRequests []map[string]json.RawMessage + client := &http.Client{Transport: brokerFixtureTransport(func(r *http.Request) (*http.Response, error) { + if !strings.HasSuffix(r.URL.Path, "/endpoint/protocols/openai/responses") { + return http.DefaultTransport.RoundTrip(r) + } + body, err := io.ReadAll(r.Body) + if err != nil { + return nil, err + } + _ = r.Body.Close() + var fields map[string]json.RawMessage + if json.Unmarshal(body, &fields) != nil { + return nil, errBrokerInvalid + } + remoteMu.Lock() + remoteRequests = append(remoteRequests, fields) + remoteMu.Unlock() + if mode == "gateway_strips_proof" { + var forwarded map[string]json.RawMessage + _ = json.Unmarshal(body, &forwarded) + delete(forwarded, "brokered_continuation_proof") + body, _ = json.Marshal(forwarded) + } + request := r.Clone(r.Context()) + request.URL, _ = url.Parse(hostedURL + "/responses") + request.Host = request.URL.Host + request.Body = io.NopCloser(bytes.NewReader(body)) + request.ContentLength = int64(len(body)) + return http.DefaultTransport.RoundTrip(request) + })} + b, brokerServer := startBrokerTestWithClient(t, cfg, client) + c := brokerTestContext(cfg) + c.LeaseExpiresAt = time.Now().Add(time.Minute).UTC().Format(time.RFC3339Nano) + proxy := httptest.NewServer(http.HandlerFunc(func(w http.ResponseWriter, r *http.Request) { + body, _ := io.ReadAll(r.Body) + if r.Header.Get("Authorization") != "Bearer fixture-acp-proxy" || + bytes.Contains(body, []byte("brokered_continuation_proof")) || bytes.Contains(body, []byte(cfg.agentKitProof)) { + t.Error("ACP request acquired broker authority") + w.WriteHeader(http.StatusBadRequest) + return + } + next := c + next.InvocationSequence = uint64(invocations.Add(1)) + next.OperationID = fmt.Sprintf("agentkit-integration-%d", next.InvocationSequence) + next.BodySHA256 = foundry.Digest(body) + raw, _ := json.Marshal(next) + r.Header.Set("Authorization", "Bearer "+brokerFixtureBearer) + r.Header.Set(brokerContextHeader, base64.RawURLEncoding.EncodeToString(raw)) + r.Body = io.NopCloser(bytes.NewReader(body)) + b.ServeHTTP(w, r) + })) + t.Cleanup(proxy.Close) + mcp := httptest.NewServer(http.HandlerFunc(func(w http.ResponseWriter, r *http.Request) { + if r.Header.Get("Authorization") != "Bearer fixture-mcp" || r.Header.Get("MCP-Protocol-Version") != "2025-06-18" { + t.Error("MCP request lost its scoped authentication") + w.WriteHeader(http.StatusUnauthorized) + return + } + var request struct { + ID json.RawMessage `json:"id"` + Method string `json:"method"` + Params struct { + Name string `json:"name"` + Arguments json.RawMessage `json:"arguments"` + } `json:"params"` + } + if json.NewDecoder(r.Body).Decode(&request) != nil { + w.WriteHeader(http.StatusBadRequest) + return + } + var result any + switch request.Method { + case "initialize": + result = map[string]any{"protocolVersion": "2025-06-18", "capabilities": map[string]any{"tools": map[string]any{}}} + case "notifications/initialized": + w.WriteHeader(http.StatusAccepted) + return + case "tools/list": + var tools []map[string]any + for _, name := range []string{"lookup_work_order", "lookup_inventory"} { + tools = append(tools, map[string]any{"name": name, "inputSchema": map[string]any{"type": "object"}}) + } + result = map[string]any{"tools": tools} + case "tools/call": + n := calls.Add(1) + if n == 1 && mode == "authorization_denied" { + w.Header().Set("Content-Type", "application/json") + _ = json.NewEncoder(w).Encode(map[string]any{"jsonrpc": "2.0", "id": request.ID, + "error": map[string]any{"code": -32001, "message": "MCP tool call is not authorized"}}) + return + } + name, arguments := "lookup_work_order", `{"id":"WO-7"}` + if n == 2 { + name, arguments = "lookup_inventory", `{"part":"FBR-7"}` + } + if n > 2 || request.Params.Name != name || string(request.Params.Arguments) != arguments { + t.Error("tool sequence changed or repeated a call") + w.WriteHeader(http.StatusBadRequest) + return + } + failed := mode == "tool_error" && n == 1 + structured := map[string]any{"part": "FBR-7", "quantity": 7} + if failed { + structured = map[string]any{"isError": true, "error": "MCP tool execution failed"} + } + text, _ := json.Marshal(structured) + result = map[string]any{"content": []map[string]string{{"type": "text", "text": string(text)}}, + "isError": failed, "structuredContent": structured, "_meta": map[string]any{"private": "not-model-visible"}} + default: + w.WriteHeader(http.StatusBadRequest) + return + } + w.Header().Set("Content-Type", "application/json") + _ = json.NewEncoder(w).Encode(map[string]any{"jsonrpc": "2.0", "id": request.ID, "result": result}) + })) + t.Cleanup(mcp.Close) + peer := brokerStartAgentKitACP(t, cfg, proxy.URL) + peer.call("initialize", map[string]any{"protocolVersion": 1, "clientCapabilities": map[string]any{}}) + cwd, err := os.Getwd() + if err != nil { + t.Fatal("could not read the ACP working directory") + } + session := peer.call("session/new", map[string]any{ + "cwd": cwd, "mcpServers": []any{map[string]any{"type": "http", "name": "governed", "url": mcp.URL + "/mcp", + "headers": []any{map[string]string{"name": "Authorization", "value": "Bearer fixture-mcp"}}}}, + }) + result, _ := session["result"].(map[string]any) + if result["sessionId"] == nil { + t.Fatal("real ACP session creation failed") + } + terminal := peer.call("session/prompt", map[string]any{"sessionId": result["sessionId"], + "prompt": []any{map[string]string{"type": "text", "text": "Check the work order, then check inventory."}}}) + blocked := mode == "authorization_denied" || mode == "gateway_strips_proof" + if blocked { + if terminal["error"] == nil || models.Load() != 1 || calls.Load() != 1 || peer.text.String() != "" { + t.Fatal("denied MCP call or missing continuation proof did not stop before model resume") + } + } else { + result, _ = terminal["result"].(map[string]any) + if result["stopReason"] != "end_turn" || models.Load() != 3 || calls.Load() != 2 || + peer.text.String() != "Checked the work order and inventory." { + t.Fatal("two governed tool rounds did not complete through the real hosted AgentKit and ACP path") + } + remoteMu.Lock() + if len(remoteRequests) != 3 || bytes.Equal(remoteRequests[1]["previous_response_id"], remoteRequests[2]["previous_response_id"]) { + t.Error("chained tool results reused a response identity") + } + remoteMu.Unlock() + } + _ = brokerTestControl(t, brokerServer.URL, brokerapi.SettlePath, c) + _ = brokerTestControl(t, brokerServer.URL, brokerapi.RetirePath, c) + for _, path := range []string{stateFile, filepath.Join(cfg.stateDir, "state.json")} { + data, err := os.ReadFile(path) + if err != nil || bytes.Contains(data, []byte(cfg.agentKitProof)) { + t.Fatal("durable state was missing or retained continuation credentials") + } + } + }) + } +} + +func brokerStartAgentKit(t *testing.T, source, modelURL string, native bool) (string, string) { + t.Helper() + source, err := filepath.Abs(source) + if err != nil { + t.Fatal("invalid AgentKit checkout path") + } + python := os.Getenv("AGENTKIT_PYTHON") + if native { + python = os.Getenv("AGENTKIT_MAF_PYTHON") + } + if python == "" { + python = "python3" + } + dir := t.TempDir() + config, stateFile := filepath.Join(dir, "agent.yaml"), filepath.Join(dir, "responses.json") + var tools []map[string]any + for _, name := range []string{"lookup_work_order", "lookup_inventory"} { + tools = append(tools, map[string]any{"name": name, "description": "Read operational data.", "brokeredClass": "read", + "parameters": map[string]any{"type": "object", "properties": map[string]any{"id": map[string]any{"type": "string"}, "part": map[string]any{"type": "string"}}}}) + } + spec := map[string]any{"abiVersion": "v0", "metadata": map[string]string{"name": "foundry-integration"}, + "model": map[string]string{"provider": "openai-compatible", "baseURL": modelURL + "/v1", "name": "fixture-model"}, + "instructions": "Use the operational tools in sequence.", "tools": []any{}, "brokeredTools": tools, "expose": map[string]any{"openai": true, "port": 8088}} + if native { + delete(spec, "brokeredTools") + spec["instructions"] = "Respond to the requested text task." + } + body, _ := json.Marshal(spec) + if os.WriteFile(config, body, 0600) != nil { + t.Fatal("could not write local AgentKit configuration") + } + listener, err := net.Listen("tcp", "127.0.0.1:0") + if err != nil { + t.Fatal("could not reserve local hosted port") + } + address := listener.Addr().String() + _, port, _ := net.SplitHostPort(address) + _ = listener.Close() + ctx, cancel := context.WithTimeout(context.Background(), time.Minute) + arguments := []string{"-m", "agentkit_serve_common.foundry_brokered_cli", "--config", config, + "--host", "127.0.0.1", "--port", port} + pythonPaths := []string{filepath.Join(source, "runtimes", "common")} + if native { + arguments = []string{"-c", `import sys, uvicorn +from agentkit_serve_common.config import load +from agentkit_serve_common.foundry import create_foundry_app +from agentkit_serve import agent_factory +app = create_foundry_app(load(sys.argv[1]), agent_factory) +uvicorn.run(app, host="127.0.0.1", port=int(sys.argv[2])) +`, config, port} + pythonPaths = append(pythonPaths, filepath.Join(source, "runtimes", "microsoft-agent-framework")) + } + command := exec.CommandContext(ctx, python, arguments...) + // No inherited model/Azure credentials, and no proof in the command line. + command.Env = []string{"PATH=" + os.Getenv("PATH"), "PYTHONPATH=" + strings.Join(pythonPaths, string(os.PathListSeparator)), + "PYTHONDONTWRITEBYTECODE=1", "AGENTKIT_FOUNDRY_BROKERED_MODEL_LOOP=1", "AGENTKIT_FOUNDRY_RESPONSE_STATE_FILE=" + stateFile, + "AGENTKIT_FOUNDRY_BROKERED_CONTINUATION_PROOF=" + brokerAgentKitFixtureProof} + var logs bytes.Buffer + command.Stdout, command.Stderr = &logs, &logs + if command.Start() != nil { + cancel() + t.Fatal("could not start AgentKit Python; set AGENTKIT_PYTHON to an environment with runtimes/common installed") + } + done := make(chan error, 1) + go func() { done <- command.Wait() }() + t.Cleanup(func() { + cancel() + <-done + if bytes.Contains(logs.Bytes(), []byte(brokerAgentKitFixtureProof)) { + t.Error("hosted AgentKit logged continuation credentials") + } + }) + base := "http://" + address + deadline := time.Now().Add(15 * time.Second) + client := &http.Client{Timeout: time.Second} + for time.Now().Before(deadline) { + response, err := client.Get(base + "/readiness") + if err == nil { + _ = response.Body.Close() + if response.StatusCode == http.StatusOK { + return base, stateFile + } + } + time.Sleep(50 * time.Millisecond) + } + t.Fatal("hosted AgentKit did not become ready; check the checkout and common-package Python dependencies") + return "", "" +} + +type brokerAgentKitACP struct { + t *testing.T + input *io.PipeWriter + decoder *json.Decoder + nextID int + text strings.Builder +} + +func brokerStartAgentKitACP(t *testing.T, cfg brokerConfiguration, proxyURL string) *brokerAgentKitACP { + t.Helper() + body, _ := json.Marshal(cfg.agent) + path := filepath.Join(t.TempDir(), "foundry.json") + if os.WriteFile(path, body, 0600) != nil { + t.Fatal("could not write local ACP configuration") + } + t.Setenv(foundry.ModelEnv, cfg.agent.Model) + t.Setenv(foundry.AgentConfigDigestEnv, foundry.Digest(body)) + t.Setenv("ORKA_FOUNDRY_ACP_PROVIDER_BASE_URL", proxyURL+"/v1") + t.Setenv("ORKA_FOUNDRY_ACP_PROVIDER_TOKEN", "fixture-acp-proxy") + input, send := io.Pipe() + receive, output := io.Pipe() + done := make(chan error, 1) + go func() { + _, err := acp.MaybeServe([]string{"--protocol", "acp", "--config", path}, input, output) + done <- err + }() + t.Cleanup(func() { + _ = send.Close() + _ = receive.Close() + select { + case <-done: + case <-time.After(5 * time.Second): + t.Error("ACP did not join after its pipes closed") + } + }) + return &brokerAgentKitACP{t: t, input: send, decoder: json.NewDecoder(receive)} +} + +func (p *brokerAgentKitACP) call(method string, params any) map[string]any { + p.t.Helper() + p.nextID++ + if json.NewEncoder(p.input).Encode(map[string]any{"jsonrpc": "2.0", "id": p.nextID, "method": method, "params": params}) != nil { + p.t.Fatal("could not write to real ACP entrypoint") + } + for { + type result struct { + message map[string]any + err error + } + done := make(chan result, 1) + go func() { + var message map[string]any + err := p.decoder.Decode(&message) + done <- result{message, err} + }() + select { + case next := <-done: + if next.err != nil { + p.t.Fatal("real ACP entrypoint closed before replying") + } + body, _ := json.Marshal(next.message) + if bytes.Contains(body, []byte(brokerAgentKitFixtureProof)) || bytes.Contains(body, []byte("caresp_")) || + bytes.Contains(body, []byte("not-model-visible")) { + p.t.Fatal("ACP output leaked continuation credentials, hosted identities, or MCP metadata") + } + if next.message["id"] == float64(p.nextID) { + return next.message + } + if next.message["method"] != "session/update" { + p.t.Fatal("ACP reply had an unexpected identity") + } + params, _ := next.message["params"].(map[string]any) + update, _ := params["update"].(map[string]any) + if update["sessionUpdate"] == "agent_message_chunk" { + content, _ := update["content"].(map[string]any) + text, _ := content["text"].(string) + p.text.WriteString(text) + } + case <-time.After(10 * time.Second): + p.t.Fatal("real ACP entrypoint did not settle") + } + } +} diff --git a/internal/broker/agentkit_test.go b/internal/broker/agentkit_test.go new file mode 100644 index 0000000..d540c28 --- /dev/null +++ b/internal/broker/agentkit_test.go @@ -0,0 +1,293 @@ +package broker + +import ( + "bytes" + "context" + "encoding/base64" + "encoding/json" + "io" + "net/http" + "net/http/httptest" + "os" + "path/filepath" + "strings" + "sync" + "testing" + + "github.com/orka-agents/agent-runtime-foundry/internal/brokerapi" + "github.com/orka-agents/agent-runtime-foundry/internal/foundry" +) + +const brokerAgentKitFixtureProof = "fixture-agentkit-continuation-proof-not-a-secret" + +func brokerAgentKitFunctionBody(previous, call, output string) []byte { + body, _ := json.Marshal(foundry.ModelResponseRequest{Model: "fixture-model", ResponseRequest: foundry.ResponseRequest{ + Stream: true, Store: true, PreviousResponseID: previous, + Input: []foundry.FunctionOutput{{Type: "function_call_output", CallID: call, Output: output}}, + }}) + return body +} + +func TestBrokerAgentKitContinuationWireAndIsolation(t *testing.T) { + for _, tc := range []struct { + name string + output string + approved bool + }{ + {"text", `{"content":[{"type":"text","text":"fixture tool result"}],"isError":false}`, true}, + {"structured", `{"content":[{"type":"text","text":"{\"count\":7}"}],"isError":false,"structuredContent":{"count":7}}`, true}, + {"execution_error", `{"content":[{"type":"text","text":"MCP tool execution failed"}],"isError":true}`, false}, + {"empty_error", `{"content":[],"isError":true}`, false}, + } { + t.Run(tc.name, func(t *testing.T) { + f := newBrokerFixture(t, "functions") + cfg := brokerTestConfig(t, f) + cfg.agentKitProof = brokerAgentKitFixtureProof + var sentMu sync.Mutex + var sent []map[string]json.RawMessage + sentCount := func() int { + sentMu.Lock() + defer sentMu.Unlock() + return len(sent) + } + client := &http.Client{Transport: brokerFixtureTransport(func(r *http.Request) (*http.Response, error) { + if strings.HasSuffix(r.URL.Path, "/responses") { + body, err := io.ReadAll(r.Body) + if err != nil { + return nil, err + } + _ = r.Body.Close() + r.Body = io.NopCloser(bytes.NewReader(body)) + var fields map[string]json.RawMessage + if json.Unmarshal(body, &fields) != nil { + t.Error("provider request was not JSON") + } + sentMu.Lock() + sent = append(sent, fields) + sentMu.Unlock() + if r.Header.Get(brokerAgentKitProofHeader) != "" { + t.Error("continuation proof was duplicated in a provider header") + } + } + return http.DefaultTransport.RoundTrip(r) + })} + _, server := startBrokerTestWithClient(t, cfg, client) + c := brokerTestContext(cfg) + status, data, err := brokerTestHTTP(context.Background(), server.URL, brokerapi.ResponsesPath, c, brokerTestBody("")) + var first foundry.Response + if err != nil || status != http.StatusOK || json.Unmarshal(data, &first) != nil || len(first.Output) != 1 { + t.Fatalf("initial tool proposal failed: status=%d", status) + } + if bytes.Contains(data, []byte(cfg.agentKitProof)) { + t.Fatal("continuation proof reached the ACP response") + } + c.InvocationSequence++ + c.OperationID = "agentkit-output" + body := brokerAgentKitFunctionBody(first.ID, first.Output[0].CallID, tc.output) + status, data, err = brokerTestHTTP(context.Background(), server.URL, brokerapi.ResponsesPath, c, body) + var final foundry.Response + if err != nil || status != http.StatusOK || json.Unmarshal(data, &final) != nil { + t.Fatalf("valid tool continuation failed: status=%d", status) + } + if bytes.Contains(data, []byte(cfg.agentKitProof)) || sentCount() != 2 { + t.Fatal("continuation escaped its provider request or was submitted more than once") + } + if _, present := sent[0]["brokered_continuation_proof"]; present { + t.Fatal("ordinary prompt carried a continuation proof") + } + var proof, previous string + if json.Unmarshal(sent[1]["brokered_continuation_proof"], &proof) != nil || proof != cfg.agentKitProof || + json.Unmarshal(sent[1]["previous_response_id"], &previous) != nil || previous != "provider-response-1" { + t.Fatal("continuation did not bind the configured proof and owned remote response") + } + var outputs []foundry.FunctionOutput + if json.Unmarshal(sent[1]["input"], &outputs) != nil || len(outputs) != 1 || outputs[0].CallID != "provider-call-id" { + t.Fatal("continuation did not translate the owned call alias") + } + var normalized struct { + Approved *bool `json:"approved"` + Output json.RawMessage `json:"output"` + Error struct { + Code string `json:"code"` + Message string `json:"message"` + } `json:"error"` + } + if json.Unmarshal([]byte(outputs[0].Output), &normalized) != nil || normalized.Approved == nil || *normalized.Approved != tc.approved { + t.Fatal("tool success/error was not preserved in the AgentKit envelope") + } + if tc.approved { + var original, forwarded any + _ = json.Unmarshal([]byte(tc.output), &original) + _ = json.Unmarshal(normalized.Output, &forwarded) + if foundry.JSONDigest(original) != foundry.JSONDigest(forwarded) || normalized.Error.Code != "" { + t.Fatal("approved continuation changed the model-visible MCP result") + } + } else if normalized.Output != nil || normalized.Error.Code != "brokered_tool_error" || normalized.Error.Message == "" { + t.Fatal("failed tool result was promoted to approval or lost its error") + } + // A duplicate must be rejected before adding a proof to another + // provider request, including a replay under a fresh operation ID. + c.InvocationSequence++ + c.OperationID = "replayed-agentkit-output" + status, _, err = brokerTestHTTP(context.Background(), server.URL, brokerapi.ResponsesPath, c, body) + if err != nil || status != http.StatusConflict || sentCount() != 2 { + t.Fatal("consumed tool output reached the provider again") + } + _ = brokerTestControl(t, server.URL, brokerapi.SettlePath, c) + // A later user prompt can continue conversation history but must + // not receive the credential reserved for tool-result continuation. + c.TaskUID, c.PromptID, c.OperationID = "next-task", "next-prompt", "next-user-prompt" + c.InvocationSequence = 1 + status, _, err = brokerTestHTTP(context.Background(), server.URL, brokerapi.ResponsesPath, c, brokerTestBody(final.ID)) + if err != nil || status != http.StatusOK || sentCount() != 3 { + t.Fatal("ordinary conversation continuation failed") + } + if _, present := sent[2]["brokered_continuation_proof"]; present { + t.Fatal("ordinary conversation continuation carried a proof") + } + for _, path := range []string{brokerapi.StatusPath, brokerapi.SettlePath, brokerapi.RetirePath} { + control := brokerTestControl(t, server.URL, path, c) + encoded, _ := json.Marshal(control) + if bytes.Contains(encoded, []byte(cfg.agentKitProof)) { + t.Fatal("continuation proof reached lifecycle status") + } + } + ledger, err := os.ReadFile(filepath.Join(cfg.stateDir, "state.json")) + if err != nil || bytes.Contains(ledger, []byte(cfg.agentKitProof)) || bytes.Contains(ledger, []byte("fixture tool result")) { + t.Fatal("ledger retained continuation credentials or tool content") + } + }) + } +} + +func TestBrokerAgentKitMalformedResultsFailBeforeReservation(t *testing.T) { + f := newBrokerFixture(t, "functions") + cfg := brokerTestConfig(t, f) + cfg.agentKitProof = brokerAgentKitFixtureProof + _, server := startBrokerTest(t, cfg) + c := brokerTestContext(cfg) + status, data, err := brokerTestHTTP(context.Background(), server.URL, brokerapi.ResponsesPath, c, brokerTestBody("")) + var first foundry.Response + if err != nil || status != http.StatusOK || json.Unmarshal(data, &first) != nil || len(first.Output) != 1 { + t.Fatal("initial proposal failed") + } + before, err := os.ReadFile(filepath.Join(cfg.stateDir, "state.json")) + if err != nil { + t.Fatal("could not read fixture ledger") + } + c.InvocationSequence++ + for name, output := range map[string]string{ + "no_error_flag": `{"content":[]}`, + "null_error_flag": `{"content":[],"isError":null}`, + "string_error_flag": `{"content":[],"isError":"false"}`, + "number_error_flag": `{"content":[],"isError":0}`, + "no_content": `{"isError":false}`, + "null_content": `{"content":null,"isError":false}`, + "unsupported_content": `{"content":[{"type":"image","text":"hidden"}],"isError":false}`, + "null_text": `{"content":[{"type":"text","text":null}],"isError":false}`, + "metadata": `{"content":[],"isError":false,"_meta":{"private":"value"}}`, + "untrusted_approval": `{"approved":true,"output":{}}`, + "duplicate_flag": `{"content":[],"isError":true,"isError":false}`, + "folded_flag": `{"content":[],"isError":true,"ISERROR":false}`, + "trailing_object": `{"content":[],"isError":false}{}`, + "invalid_unicode": `{"content":[{"type":"text","text":"\ud800"}],"isError":false}`, + "array_structured": `{"content":[],"isError":false,"structuredContent":[]}`, + "null_structured": `{"content":[],"isError":false,"structuredContent":null}`, + } { + t.Run(name, func(t *testing.T) { + c.OperationID = "invalid-" + name + body := brokerAgentKitFunctionBody(first.ID, first.Output[0].CallID, output) + status, _, err := brokerTestHTTP(context.Background(), server.URL, brokerapi.ResponsesPath, c, body) + if err != nil || status != http.StatusBadRequest { + t.Fatalf("invalid MCP result was accepted: status=%d", status) + } + }) + } + after, err := os.ReadFile(filepath.Join(cfg.stateDir, "state.json")) + _, inferences, _, _ := f.counts() + if err != nil || !bytes.Equal(before, after) || inferences != 1 { + t.Fatal("malformed tool output changed durable ownership or reached the provider") + } +} + +func TestBrokerRejectsChildSuppliedAgentKitProof(t *testing.T) { + for _, configured := range []bool{false, true} { + for _, where := range []string{"body", "folded_body", "empty_body", "header", "empty_header", "folded_header"} { + t.Run(where+"/"+map[bool]string{false: "disabled", true: "enabled"}[configured], func(t *testing.T) { + f := newBrokerFixture(t, "success") + cfg := brokerTestConfig(t, f) + if configured { + cfg.agentKitProof = brokerAgentKitFixtureProof + } + b, _ := startBrokerTest(t, cfg) + body := brokerTestBody("") + switch where { + case "body": + body = append([]byte(`{"brokered_continuation_proof":"child-controlled",`), body[1:]...) + case "folded_body": + body = append([]byte(`{"Brokered_Continuation_Proof":"child-controlled",`), body[1:]...) + case "empty_body": + body = append([]byte(`{"brokered_continuation_proof":null,`), body[1:]...) + } + c := brokerTestContext(cfg) + c.BodySHA256 = foundry.Digest(body) + raw, _ := json.Marshal(c) + r := httptest.NewRequest(http.MethodPost, brokerapi.ResponsesPath, bytes.NewReader(body)) + r.Header.Set("Authorization", "Bearer "+brokerFixtureBearer) + r.Header.Set(brokerContextHeader, base64.RawURLEncoding.EncodeToString(raw)) + switch where { + case "header": + r.Header.Set(brokerAgentKitProofHeader, "child-controlled") + case "empty_header": + r.Header.Set(brokerAgentKitProofHeader, "") + case "folded_header": + r.Header["x-AgEnTkIt-brokered-continuation-proof"] = []string{"child-controlled"} + } + w := httptest.NewRecorder() + b.ServeHTTP(w, r) + if w.Code != http.StatusBadRequest || len(b.ledger.Sessions) != 0 { + t.Fatal("child-controlled proof was admitted or reserved ownership") + } + creates, inferences, _, _ := f.counts() + if creates+inferences != 0 || bytes.Contains(w.Body.Bytes(), []byte("child-controlled")) { + t.Fatal("rejected proof reached the provider or error response") + } + }) + } + } +} + +func TestBrokerAgentKitProofConfiguration(t *testing.T) { + f := newBrokerFixture(t, "success") + cfg := brokerTestConfig(t, f) + body, _ := json.Marshal(cfg.agent) + path := filepath.Join(t.TempDir(), "foundry.json") + if os.WriteFile(path, body, 0600) != nil { + t.Fatal("could not write fixture config") + } + env := map[string]string{ + foundry.ModelEnv: cfg.agent.Model, foundry.AgentConfigDigestEnv: foundry.Digest(body), + "ORKA_FOUNDRY_BROKER_STATE_DIR": cfg.stateDir, "ORKA_FOUNDRY_BROKER_BEARER_TOKEN": cfg.bearer, + } + for _, tc := range []struct { + name string + proof string + valid bool + }{ + {"disabled", "", true}, {"configured", brokerAgentKitFixtureProof, true}, + {"short", "short", false}, {"space", strings.Repeat("x", 32) + " ", false}, + {"unicode_space", strings.Repeat("x", 32) + "\u00a0", false}, {"invalid_utf8", strings.Repeat("x", 32) + "\xff", false}, + {"control", strings.Repeat("x", 32) + "\n", false}, {"oversized", strings.Repeat("x", (16<<10)+1), false}, + } { + t.Run(tc.name, func(t *testing.T) { + env[brokerAgentKitProofEnv] = tc.proof + loaded, err := loadBrokerConfiguration(path, func(name string) string { return env[name] }) + if (err == nil) != tc.valid || (err == nil && loaded.agentKitProof != tc.proof) { + t.Fatal("continuation proof configuration accepted an invalid value or lost the configured value") + } + if err != nil && tc.proof != "" && strings.Contains(err.Error(), tc.proof) { + t.Fatal("configuration failure disclosed the proof") + } + }) + } +} diff --git a/internal/broker/broker.go b/internal/broker/broker.go index 8200d22..4fc5a99 100644 --- a/internal/broker/broker.go +++ b/internal/broker/broker.go @@ -7,6 +7,7 @@ import ( "encoding/json" "errors" "io" + "log/slog" "net/http" "sync" "time" @@ -42,10 +43,11 @@ type lifecycleBroker struct { ctx context.Context cancel context.CancelFunc wg sync.WaitGroup + diagnosticLog *slog.Logger } func newLifecycleBroker(ctx context.Context, cfg brokerConfiguration, provider foundry.TokenProvider, client *http.Client) (*lifecycleBroker, error) { - if provider == nil || len(cfg.bearer) < 32 || !foundry.DigestValid(cfg.configDigest) { + if provider == nil || len(cfg.bearer) < 32 || !foundry.DigestValid(cfg.configDigest) || !brokerAgentKitProofValid(cfg.agentKitProof) { return nil, errBrokerInvalid } store, ledger, err := openBrokerStore(cfg.stateDir, cfg.configDigest) diff --git a/internal/broker/response_diagnostics.go b/internal/broker/response_diagnostics.go new file mode 100644 index 0000000..537f295 --- /dev/null +++ b/internal/broker/response_diagnostics.go @@ -0,0 +1,100 @@ +package broker + +import ( + "encoding/json" + + "github.com/orka-agents/agent-runtime-foundry/internal/foundry" + "github.com/orka-agents/agent-runtime-foundry/internal/strictjson" +) + +// Diagnostics never participate in ownership, output validation, or settlement. +// One bounded record is emitted per rejected invocation. Provider text and IDs +// are deliberately absent, including from errors and unknown metadata values. +type brokerResponseDiagnostic struct { + stage string + httpStatus int + accepted bool + streamComplete bool + terminalFrames int + terminalStatus string + errorCode string + upstreamStatus int +} + +func (d *brokerResponseDiagnostic) observe(raw []byte, response foundry.Response) { + if d == nil { + return + } + d.accepted = true + switch response.Status { + case "completed", "failed", "incomplete", "cancelled": + default: + return + } + d.terminalFrames++ + if d.terminalFrames != 1 { + d.terminalStatus, d.errorCode, d.upstreamStatus = "multiple", "", 0 + return + } + d.terminalStatus = response.Status + if response.Error == nil { + return + } + d.errorCode = brokerSafeResponseCode(response.Error.Code) + if d.errorCode == "unknown" { + return + } + // This optional metadata is decoded separately so it cannot change the + // acceptance of an otherwise valid response or the existing wire errors. + var envelope struct { + Error json.RawMessage `json:"error"` + } + var metadata struct { + UpstreamStatus json.RawMessage `json:"upstream_status"` + } + var status int + if strictjson.DecodeStruct(raw, &envelope, false) == nil && + strictjson.DecodeStruct(envelope.Error, &metadata, false) == nil && + json.Unmarshal(metadata.UpstreamStatus, &status) == nil && status >= 400 && status <= 599 { + d.upstreamStatus = status + } +} + +func brokerSafeResponseCode(code string) string { + switch code { + case "ModelAuthMissing", "ModelAuthRejected", "ModelUnavailable", "ModelUpstreamError", + "InvalidModelResponse", "ModelResponseTooLarge", "ModelResumeError", + "tool_loop_limit_exceeded", "brokered_response_state_storage_error", + "brokered_response_state_too_large", "brokered_response_state_full": + return code + default: + return "unknown" + } +} + +func (b *lifecycleBroker) logResponseFailure(c brokerContext, d brokerResponseDiagnostic) { + if b.diagnosticLog == nil { + return + } + fields := []any{ + "stage", d.stage, + "owner_digest", foundry.JSONDigest(c.Owner), + "invocation_sequence", c.InvocationSequence, + "response_acknowledged", d.accepted, + "stream_complete", d.streamComplete, + "terminal_frames", d.terminalFrames, + } + if d.httpStatus >= 100 && d.httpStatus <= 599 { + fields = append(fields, "http_status", d.httpStatus) + } + if d.terminalStatus != "" { + fields = append(fields, "observed_terminal_status", d.terminalStatus) + } + if d.errorCode != "" { + fields = append(fields, "error_code", d.errorCode) + } + if d.upstreamStatus != 0 { + fields = append(fields, "upstream_status", d.upstreamStatus) + } + b.diagnosticLog.Warn("Foundry response rejected", fields...) +} diff --git a/internal/broker/response_diagnostics_test.go b/internal/broker/response_diagnostics_test.go new file mode 100644 index 0000000..5f397db --- /dev/null +++ b/internal/broker/response_diagnostics_test.go @@ -0,0 +1,260 @@ +package broker + +import ( + "bytes" + "context" + "encoding/json" + "errors" + "io" + "log/slog" + "net/http" + "net/http/httptest" + "os" + "strings" + "testing" + + "github.com/orka-agents/agent-runtime-foundry/internal/brokerapi" + "github.com/orka-agents/agent-runtime-foundry/internal/foundry" +) + +func TestBrokerResponseFailureDiagnostics(t *testing.T) { + for _, test := range []struct { + name string + stage string + terminal string + code string + upstream int + frames int + accepted bool + complete bool + successful bool + }{ + {name: "failed", stage: "strict_decode", terminal: "failed", code: "ModelAuthRejected", upstream: 403, frames: 1, accepted: true, complete: true}, + {name: "json-failed", stage: "response_terminal", terminal: "failed", code: "ModelAuthRejected", upstream: 403, frames: 1, accepted: true}, + {name: "unknown-code", stage: "strict_decode", terminal: "failed", code: "unknown", frames: 1, accepted: true, complete: true}, + {name: "incomplete", stage: "strict_decode", terminal: "incomplete", frames: 1, accepted: true, complete: true}, + {name: "cancelled", stage: "strict_decode", terminal: "cancelled", frames: 1, accepted: true, complete: true}, + {name: "malformed-success", stage: "strict_decode", terminal: "completed", frames: 1, accepted: true, complete: true}, + {name: "wrong-identity", stage: "stream_read"}, + {name: "missing-terminal", stage: "strict_decode", accepted: true, complete: true}, + {name: "malformed-tail", stage: "stream_read", accepted: true}, + {name: "multiple-terminals", stage: "strict_decode", terminal: "multiple", frames: 2, accepted: true, complete: true}, + {name: "success", successful: true}, + } { + t.Run(test.name, func(t *testing.T) { + f := newBrokerEvidenceFixture(t, func(request foundry.ResponseRequest) (string, []byte) { + response := map[string]any{ + "id": "provider-response-do-not-log", "status": "in_progress", + "agent_session_id": request.AgentSessionID, "output": []any{}, + } + frame := func(kind string) string { + raw, _ := json.Marshal(map[string]any{"type": kind, "response": response}) + return testSSE(string(raw)) + } + created := frame("response.created") + response["status"] = "failed" + response["error"] = map[string]any{ + "code": "ModelAuthRejected", "message": "provider-message-do-not-log", + "upstream_status": 403, "request_id": "provider-request-do-not-log", + } + switch test.name { + case "json-failed": + raw, _ := json.Marshal(response) + return "application/json", raw + case "unknown-code": + response["error"].(map[string]any)["code"] = "Bearer credential-shaped-code-do-not-log" + case "incomplete": + delete(response, "error") + response["status"] = "incomplete" + response["incomplete_details"] = map[string]any{"reason": "provider-reason-do-not-log"} + case "cancelled": + delete(response, "error") + response["status"] = "cancelled" + case "malformed-success": + delete(response, "error") + response["status"] = "completed" + response["output"] = []any{map[string]any{"type": "untrusted_native_tool", "name": "provider-native-tool-do-not-log"}} + case "wrong-identity": + response["agent_session_id"] = "provider-unowned-session-do-not-log" + return "text/event-stream", []byte(frame("response.failed")) + case "missing-terminal": + return "text/event-stream", []byte(created) + case "malformed-tail": + return "text/event-stream", []byte(created + "data: {\n\n") + case "multiple-terminals": + return "text/event-stream", []byte(created + frame("response.failed") + frame("response.failed")) + case "success": + delete(response, "error") + response["status"] = "completed" + response["output"] = []any{map[string]any{"type": "message", "role": "assistant", "content": []any{map[string]any{"type": "output_text", "text": "provider-success-do-not-log"}}}} + } + return "text/event-stream", []byte(created + frame("response."+response["status"].(string))) + }) + cfg := brokerTestConfig(t, f) + b, server := startBrokerTest(t, cfg) + var logs bytes.Buffer + b.diagnosticLog = slog.New(slog.NewJSONHandler(&logs, nil)) + c := brokerTestContext(cfg) + status, body, err := brokerTestHTTP(context.Background(), server.URL, brokerapi.ResponsesPath, c, brokerTestBody("")) + if err != nil || (status == http.StatusOK) != test.successful { + t.Fatal("diagnostics changed inference outcome") + } + if test.successful { + if logs.Len() != 0 { + t.Fatal("successful inference emitted a failure diagnostic") + } + return + } + record := decodeBrokerDiagnostic(t, logs.Bytes()) + if record["stage"] != test.stage || record["response_acknowledged"] != test.accepted || + record["stream_complete"] != test.complete || record["terminal_frames"] != float64(test.frames) || + record["http_status"] != float64(200) || record["invocation_sequence"] != float64(1) || + record["owner_digest"] != foundry.JSONDigest(c.Owner) { + t.Fatal("diagnostic did not identify the bounded failure boundary") + } + if terminal, _ := record["observed_terminal_status"].(string); terminal != test.terminal { + t.Fatal("diagnostic conflated provider failure and malformed completion") + } + if code, _ := record["error_code"].(string); code != test.code { + t.Fatal("diagnostic did not preserve only the allowlisted error code") + } + if upstream, _ := record["upstream_status"].(float64); upstream != float64(test.upstream) { + t.Fatal("diagnostic did not bound the upstream status") + } + for _, forbidden := range []string{ + "do-not-log", brokerFixtureBearer, brokerTestToken(), "fixture-input-do-not-persist", + c.Owner.RuntimeSessionUID, c.TaskUID, c.PromptID, f.server.URL, + } { + if bytes.Contains(logs.Bytes(), []byte(forbidden)) { + t.Fatal("diagnostic exposed provider or owner data") + } + } + if bytes.Contains(body, []byte("ModelAuthRejected")) || bytes.Contains(body, []byte("upstream_status")) { + t.Fatal("diagnostics changed the child error envelope") + } + if !test.accepted { + if brokerInvocationState(b, c) != "uncertain" { + t.Fatal("diagnostics fabricated remote acknowledgement") + } + brokerPendingControl(t, server.URL, brokerapi.SettlePath, c, false, 1) + return + } + proof := brokerTestControl(t, server.URL, brokerapi.SettlePath, c) + if !proof.SettlementProven || proof.AmbiguousInvocations != 0 { + t.Fatal("diagnostics changed acknowledged failure settlement") + } + proof = brokerTestControl(t, server.URL, brokerapi.RetirePath, c) + if !proof.RetirementProven { + t.Fatal("diagnostics changed owner retirement") + } + _ = decodeBrokerDiagnostic(t, logs.Bytes()) + }) + } +} + +func decodeBrokerDiagnostic(t *testing.T, data []byte) map[string]any { + t.Helper() + var record map[string]any + decoder := json.NewDecoder(bytes.NewReader(data)) + if decoder.Decode(&record) != nil { + t.Fatal("failure did not emit a JSON diagnostic") + } + var extra any + if !errors.Is(decoder.Decode(&extra), io.EOF) { + t.Fatal("one rejected invocation emitted multiple diagnostics") + } + allowed := map[string]bool{} + for _, field := range []string{"time", "level", "msg", "stage", "owner_digest", "invocation_sequence", "response_acknowledged", "stream_complete", "terminal_frames", "http_status", "observed_terminal_status", "error_code", "upstream_status"} { + allowed[field] = true + } + for field := range record { + if !allowed[field] { + t.Fatal("diagnostic emitted an unreviewed field") + } + } + return record +} + +func TestBrokerResponseDiagnosticRejectsUnsafeStatusMetadata(t *testing.T) { + for _, raw := range []string{`"403"`, `403.0`, `true`, `null`, `399`, `600`, `{"secret":"do-not-log"}`} { + t.Run(raw, func(t *testing.T) { + diagnostic := brokerResponseDiagnostic{} + data := []byte(`{"error":{"code":"ModelAuthRejected","message":"do-not-log","upstream_status":` + raw + `}}`) + diagnostic.observe(data, foundry.Response{Status: "failed", Error: &foundry.ResponseError{Code: "ModelAuthRejected"}}) + if diagnostic.upstreamStatus != 0 || diagnostic.errorCode != "ModelAuthRejected" { + t.Fatal("untrusted status metadata entered the diagnostic") + } + }) + } +} + +func TestBrokerResponseDiagnosticDoesNotClaimFailedIdentityWrite(t *testing.T) { + b, c := newBrokerResponseIdentityFixture(t) + before := brokerIdentityLedgerBytes(t, b) + if err := os.Rename(b.store.dir, b.store.dir+"-unavailable"); err != nil { + t.Fatal("could not inject identity persistence failure") + } + diagnostic := brokerResponseDiagnostic{} + frame := testSSE(`{"type":"response.failed","response":{"id":"provider-do-not-log","status":"failed","error":{"code":"ModelAuthRejected","message":"do-not-log","upstream_status":403}}}`) + _, err := b.readTrackedStream(strings.NewReader(frame), c, brokerIdentityRemoteSession, &diagnostic) + if !errors.Is(err, errBrokerStorage) || diagnostic.accepted || diagnostic.errorCode != "" || diagnostic.terminalFrames != 0 { + t.Fatal("diagnostics claimed acceptance after failed persistence") + } + if foundry.Digest(before) != foundry.JSONDigest(b.ledger) { + t.Fatal("diagnostics changed failed-write authority") + } +} + +func TestBrokerResponseDiagnosticDistinguishesHTTPAndTransportFailure(t *testing.T) { + for _, test := range []struct { + name string + stage string + httpStatus int + state string + }{ + {name: "rejected", stage: "http_response", httpStatus: 429, state: "rejected"}, + {name: "rejected-server", stage: "http_response", httpStatus: 503, state: "uncertain"}, + {name: "transport", stage: "transport", state: "uncertain"}, + } { + t.Run(test.name, func(t *testing.T) { + f := newBrokerFixture(t, test.name) + if test.name == "transport" { + f.server.Close() + f.server = httptest.NewServer(http.HandlerFunc(func(w http.ResponseWriter, r *http.Request) { + if r.Method != http.MethodPost || !strings.HasSuffix(r.URL.Path, "/protocols/openai/responses") { + f.serve(w, r) + return + } + _, _ = io.Copy(io.Discard, r.Body) + connection, _, err := w.(http.Hijacker).Hijack() + if err != nil { + t.Error("could not inject response transport loss") + return + } + _ = connection.Close() + })) + } + cfg := brokerTestConfig(t, f) + b, server := startBrokerTest(t, cfg) + var logs bytes.Buffer + b.diagnosticLog = slog.New(slog.NewJSONHandler(&logs, nil)) + c := brokerTestContext(cfg) + status, _, err := brokerTestHTTP(context.Background(), server.URL, brokerapi.ResponsesPath, c, brokerTestBody("")) + state := brokerInvocationState(b, c) + validState := state == test.state || (test.state == "rejected" && state == "settled") + if err != nil || status == http.StatusOK || !validState { + t.Fatal("diagnostics changed rejected or uncertain ownership") + } + record := decodeBrokerDiagnostic(t, logs.Bytes()) + httpStatus, _ := record["http_status"].(float64) + if record["stage"] != test.stage || httpStatus != float64(test.httpStatus) || record["response_acknowledged"] != false { + t.Fatal("diagnostic conflated HTTP rejection and transport loss") + } + if test.state == "uncertain" { + brokerPendingControl(t, server.URL, brokerapi.SettlePath, c, false, 1) + } else if !brokerTestControl(t, server.URL, brokerapi.SettlePath, c).SettlementProven { + t.Fatal("definite HTTP rejection could not settle") + } + }) + } +} diff --git a/internal/broker/response_identity_write_test.go b/internal/broker/response_identity_write_test.go index 865b46f..9996817 100644 --- a/internal/broker/response_identity_write_test.go +++ b/internal/broker/response_identity_write_test.go @@ -131,7 +131,7 @@ func TestBrokerResponseIdentityWritesOncePerInvocation(t *testing.T) { events[0] = testSSE(`{"type":"response.created","response":{"id":"response-1","status":"in_progress"}}`) events[len(events)-1] = testSSE(`{"type":"response.completed","response":{"id":"response-1","status":"completed"}}`) reader := &brokerIdentityEventReader{t: t, path: filepath.Join(b.store.dir, "state.json"), events: events} - data, err := b.readTrackedStream(reader, c, brokerIdentityRemoteSession) + data, err := b.readTrackedStream(reader, c, brokerIdentityRemoteSession, nil) if err != nil { t.Fatal("coherent acceptance evidence was rejected") } @@ -287,7 +287,7 @@ func TestBrokerResponseIdentityTrackingStillRejectsMalformedTail(t *testing.T) { b, c := newBrokerResponseIdentityFixture(t) created := testSSE(`{"type":"response.created","response":{"id":"response-1","status":"in_progress"}}`) stream := created + testSSE(`{"type":"response.in_progress","response":{"id":"response-1","status":"completed"}}`) - if _, err := b.readTrackedStream(strings.NewReader(stream), c, brokerIdentityRemoteSession); err == nil { + if _, err := b.readTrackedStream(strings.NewReader(stream), c, brokerIdentityRemoteSession, nil); err == nil { t.Fatal("duplicate response identity bypassed lifecycle validation") } invocation := b.ledger.Sessions[foundry.JSONDigest(c.Owner)].Prompts[c.promptKey()].Invocations[c.InvocationSequence] diff --git a/internal/broker/responses.go b/internal/broker/responses.go index 406ff07..e2c3b90 100644 --- a/internal/broker/responses.go +++ b/internal/broker/responses.go @@ -19,6 +19,12 @@ import ( ) func (b *lifecycleBroker) serveResponses(w http.ResponseWriter, r *http.Request, c brokerContext, raw []byte) { + for name := range r.Header { + if strings.EqualFold(name, brokerAgentKitProofHeader) { + brokerWriteError(w, errBrokerInvalid) + return + } + } var request foundry.ModelResponseRequest if strictjson.Decode(raw, &request, true) != nil || request.Model != b.cfg.agent.Model || !request.Stream || !request.Store || request.Input == nil { brokerWriteError(w, errBrokerInvalid) @@ -128,6 +134,11 @@ func (b *lifecycleBroker) reserveInvocation(c brokerContext, request *foundry.Mo if err := brokerTranslatePrevious(session, c, prompt.LastSequence == 0, &request.ResponseRequest); err != nil { return err } + if b.cfg.agentKitProof != "" { + if err := brokerAgentKitOutputs(&request.ResponseRequest); err != nil { + return err + } + } prompt.LastSequence = c.InvocationSequence prompt.Invocations[c.InvocationSequence] = &brokerInvocation{Sequence: c.InvocationSequence, OperationID: c.OperationID, BodyDigest: c.BodySHA256, State: "reserved"} @@ -184,7 +195,7 @@ func brokerTranslatePrevious(session *brokerSession, c brokerContext, first bool return nil } -func (b *lifecycleBroker) invoke(ctx context.Context, c brokerContext, request foundry.ResponseRequest) (foundry.Response, error) { +func (b *lifecycleBroker) invoke(ctx context.Context, c brokerContext, request foundry.ResponseRequest) (_ foundry.Response, resultErr error) { key, promptKey := foundry.JSONDigest(c.Owner), c.promptKey() b.mu.Lock() session := b.ledger.Sessions[key] @@ -263,7 +274,7 @@ func (b *lifecycleBroker) invoke(ctx context.Context, c brokerContext, request f return foundry.Response{}, errBrokerRemote } request.AgentSessionID = remoteID - body, err := json.Marshal(request) + body, err := b.marshalResponseRequest(request) if err != nil { return foundry.Response{}, errBrokerInvalid } @@ -285,12 +296,20 @@ func (b *lifecycleBroker) invoke(ctx context.Context, c brokerContext, request f if err != nil { return foundry.Response{}, err } + diagnostic := brokerResponseDiagnostic{stage: "transport"} + defer func() { + if resultErr != nil { + b.logResponseFailure(c, diagnostic) + } + }() response, err := b.sendRemoteRequest(prepared) if err != nil { return foundry.Response{}, err } defer response.Body.Close() //nolint:errcheck + diagnostic.httpStatus = response.StatusCode if response.StatusCode != http.StatusOK { + diagnostic.stage = "http_response" count, readErr := io.Copy(io.Discard, io.LimitReader(response.Body, foundry.MaxAgentConfigBytes+1)) if readErr != nil || count > foundry.MaxAgentConfigBytes || !foundry.DefiniteRejection(response.StatusCode) { return foundry.Response{}, errBrokerAmbiguous @@ -308,6 +327,7 @@ func (b *lifecycleBroker) invoke(ctx context.Context, c brokerContext, request f } return foundry.Response{}, errBrokerRemote } + diagnostic.stage = "content_type" mediaType, _, err := mime.ParseMediaType(response.Header.Get("Content-Type")) if err != nil { return foundry.Response{}, errBrokerAmbiguous @@ -316,25 +336,36 @@ func (b *lifecycleBroker) invoke(ctx context.Context, c brokerContext, request f switch mediaType { case "text/event-stream": var data []byte - data, err = b.readTrackedStream(response.Body, c, remoteID) + diagnostic.stage = "stream_read" + data, err = b.readTrackedStream(response.Body, c, remoteID, &diagnostic) if err == nil { + diagnostic.stage = "strict_decode" + diagnostic.streamComplete = true summary, err = foundry.ParseStrictSSE(bytes.NewReader(data)) } case "application/json": var data []byte + diagnostic.stage = "response_read" data, err = io.ReadAll(io.LimitReader(response.Body, foundry.DefaultMaxStreamBytes+1)) if err == nil && len(data) <= foundry.DefaultMaxStreamBytes { var document foundry.Response + diagnostic.stage = "response_decode" document, err = brokerDecodeResponseEvidence(data) if err == nil { + diagnostic.stage = "response_identity" err = b.recordResponseIdentity(c, document, remoteID) + if err == nil { + diagnostic.observe(data, document) + } } if err == nil && document.Status == "completed" { + diagnostic.stage = "strict_decode" document, err = foundry.DecodeResponse(data) if err == nil { summary, err = foundry.CompleteResponse(document, foundry.ResponseCallbacks{}) } } else if err == nil { + diagnostic.stage = "response_terminal" err = errBrokerRemote } } else { @@ -346,6 +377,7 @@ func (b *lifecycleBroker) invoke(ctx context.Context, c brokerContext, request f if err != nil || foundry.ValidateSummary(summary) != nil { return foundry.Response{}, errBrokerAmbiguous } + diagnostic.stage = "completion_storage" return b.commitCompletedResponse(c, summary) } @@ -385,7 +417,7 @@ func brokerDecodeResponseEvidence(data []byte) (foundry.Response, error) { return response, nil } -func (b *lifecycleBroker) readTrackedStream(reader io.Reader, c brokerContext, remoteID string) ([]byte, error) { +func (b *lifecycleBroker) readTrackedStream(reader io.Reader, c brokerContext, remoteID string, diagnostic *brokerResponseDiagnostic) ([]byte, error) { limited := &io.LimitedReader{R: reader, N: foundry.DefaultMaxStreamBytes + 1} scanner := bufio.NewScanner(limited) scanner.Buffer(make([]byte, 32<<10), foundry.DefaultMaxEventBytes) @@ -421,7 +453,11 @@ func (b *lifecycleBroker) readTrackedStream(reader io.Reader, c brokerContext, r default: return errBrokerRemote } - return b.recordResponseIdentity(c, response, remoteID) + if err := b.recordResponseIdentity(c, response, remoteID); err != nil { + return err + } + diagnostic.observe(frame.Response, response) + return nil } for scanner.Scan() { line := scanner.Bytes() diff --git a/internal/broker/run.go b/internal/broker/run.go index 4e1c275..a233cfc 100644 --- a/internal/broker/run.go +++ b/internal/broker/run.go @@ -5,6 +5,7 @@ import ( "errors" "flag" "io" + "log/slog" "net" "net/http" "os" @@ -23,6 +24,7 @@ type brokerConfiguration struct { addr string stateDir string bearer string + agentKitProof string operationTimeout time.Duration } @@ -62,6 +64,7 @@ func MaybeServe(args []string) (bool, error) { return true, err } defer broker.close() + broker.diagnosticLog = slog.New(slog.NewJSONHandler(os.Stderr, nil)) server := &http.Server{Addr: cfg.addr, Handler: broker, ReadHeaderTimeout: 5 * time.Second, ReadTimeout: 30 * time.Second, IdleTimeout: 30 * time.Second, MaxHeaderBytes: 32 << 10} finished := make(chan error, 1) @@ -98,9 +101,11 @@ func loadBrokerConfiguration(path string, getenv func(string) string) (brokerCon cfg := brokerConfiguration{agent: agent, configDigest: digest, addr: foundry.FirstNonBlank(getenv("ORKA_FOUNDRY_BROKER_ADDR"), "127.0.0.1:8091"), stateDir: getenv("ORKA_FOUNDRY_BROKER_STATE_DIR"), bearer: getenv("ORKA_FOUNDRY_BROKER_BEARER_TOKEN"), + agentKitProof: getenv(brokerAgentKitProofEnv), operationTimeout: 45 * time.Second} if !brokerAddressValid(cfg.addr) || cfg.stateDir == "" || !foundry.SafeString(cfg.bearer, 16<<10) || len(cfg.bearer) < 32 || strings.ContainsAny(cfg.bearer, " \t") || + !brokerAgentKitProofValid(cfg.agentKitProof) || (getenv(foundry.IsolationModeEnv) != "" && getenv(foundry.IsolationModeEnv) != "entra") { return brokerConfiguration{}, errBrokerInvalid }