diff --git a/Dockerfile.acp b/Dockerfile.acp new file mode 100644 index 0000000..f88a9d3 --- /dev/null +++ b/Dockerfile.acp @@ -0,0 +1,20 @@ +# syntax=docker/dockerfile:1 +FROM --platform=$BUILDPLATFORM golang:1.26.5-alpine AS build +ARG TARGETOS +ARG TARGETARCH +WORKDIR /src +COPY go.mod go.sum ./ +RUN go mod download +COPY . . +RUN CGO_ENABLED=0 GOOS=$TARGETOS GOARCH=$TARGETARCH \ + go build -buildvcs=false -trimpath -ldflags='-s -w' -o /out/agent-runtime-foundry . +ARG FOUNDRY_CONFIG=examples/foundry-acp.json +COPY ${FOUNDRY_CONFIG} /out/rootfs/agent/foundry.json +RUN chmod 0555 /out/rootfs/agent && chmod 0444 /out/rootfs/agent/foundry.json + +FROM gcr.io/distroless/static:nonroot@sha256:1c2c046bc09ed40fad370b599a0b1ae7987f55b01e247cf27a7c27cd97e5bbc7 +COPY --from=build --chown=0:0 --chmod=0555 /out/agent-runtime-foundry /agent-runtime-foundry +COPY --from=build --chown=0:0 /out/rootfs/ / +USER 65532:65532 +ENTRYPOINT ["/agent-runtime-foundry"] +CMD ["--protocol", "acp", "--config", "/agent/foundry.json"] diff --git a/Dockerfile.hosted b/Dockerfile.hosted new file mode 100644 index 0000000..466a38b --- /dev/null +++ b/Dockerfile.hosted @@ -0,0 +1,39 @@ +# syntax=docker/dockerfile:1 +# Build the configured ACP image, compose it with Orka's supervisor, then use +# that immutable composition here. Only the public hosted configuration enters +# this image. The gateway retains Azure identity and all bootstrap secrets. +ARG ORKA_RUNTIME_IMAGE +FROM ${ORKA_RUNTIME_IMAGE} AS runtime + +# This filesystem was exercised on Foundry Hosted Agents with the actual Orka +# exec helper. A distroless filesystem did not start in that environment. +FROM docker.io/library/python:3.12-slim@sha256:2fe5997d249a808b8eeea52c58a1dbffbba28754dc11699ef5c029f2d818ce79 +ARG ORKA_RUNTIME_IMAGE +ARG FOUNDRY_ADAPTER_DIGEST +ARG HOSTED_CONFIG +RUN set -eu; \ + case "$ORKA_RUNTIME_IMAGE" in *@sha256:*) ;; *) exit 1 ;; esac; \ + digest="${ORKA_RUNTIME_IMAGE##*@sha256:}"; \ + test "${#digest}" -eq 64; \ + case "$digest" in *[!0-9a-f]*) exit 1 ;; esac; \ + case "$FOUNDRY_ADAPTER_DIGEST" in sha256:*) ;; *) exit 1 ;; esac; \ + digest="${FOUNDRY_ADAPTER_DIGEST#sha256:}"; \ + test "${#digest}" -eq 64; \ + case "$digest" in *[!0-9a-f]*) exit 1 ;; esac; \ + mkdir -p /agent; chmod 0555 /agent +COPY --from=runtime --chown=0:0 --chmod=0555 /agent-runtime-foundry /agent-runtime-foundry +COPY --from=runtime --chown=0:0 --chmod=0555 /usr/local/bin/orka-acp-runtime /usr/local/bin/orka-acp-runtime +COPY --from=runtime --chown=0:0 --chmod=0555 /usr/local/bin/orka-acp-exec-helper /usr/local/bin/orka-acp-exec-helper +COPY --from=runtime --chown=0:0 --chmod=0444 /agent/foundry.json /agent/foundry.json +COPY --chown=0:0 --chmod=0444 ${HOSTED_CONFIG} /agent/hosted.json +ENV PORT=8088 ORKA_ACP_FOUNDRY_ADAPTER_DIGEST=${FOUNDRY_ADAPTER_DIGEST} +LABEL org.opencontainers.image.title="Orka-governed Foundry Hosted Agent" \ + org.opencontainers.image.source="https://github.com/orka-agents/agent-runtime-foundry" \ + io.orka.harness.protocol="orka.harness.v2" +WORKDIR / +# Only the supervisor is privileged. It assigns distinct UIDs/GIDs to ACP children. +USER 0:0 +EXPOSE 8088 +STOPSIGNAL SIGTERM +ENTRYPOINT ["/agent-runtime-foundry"] +CMD ["--protocol", "hosted", "--config", "/agent/hosted.json"] diff --git a/README.md b/README.md index 17dc0be..da8723b 100644 --- a/README.md +++ b/README.md @@ -1,10 +1,15 @@ # Microsoft Foundry Hosted Agents adapter for Orka -This repository presents a deployed **Microsoft Foundry Hosted Agent** that -implements the Responses protocol as an +This repository connects a deployed **Microsoft Foundry Hosted Agent** to Orka. +For harness v2, use the [ACP child and durable Foundry broker](docs/harness-v2.md) +with Orka's existing supervisor. The default HTTP entry point exposes the +Responses agent as an [`orka.harness.v1`](https://github.com/orka-agents/orka/blob/main/website/docs/development/agent-runtime-adapter-contract.md) `AgentRuntime` endpoint. +To run the supervisor and ACP child inside Foundry itself, use the +[Hosted Agent v2 package and Kubernetes gateway](docs/foundry-hosted-v2.md). + The adapter calls the Hosted Agent's dedicated Responses endpoint: ```text @@ -21,7 +26,7 @@ are never sent to Foundry. ## Status -The adapter is experimental. Run a single replica. Runtime-session and active +The adapter is experimental. Run a single replica. In harness v1, runtime-session and active turn state are currently process-local, so a pod replacement cannot resume a retained session or deduplicate an active turn. Orka facade samples use an external endpoint and do not install or manage this adapter. @@ -31,15 +36,20 @@ not install or manage this adapter. Deploy a Hosted Agent that exposes the Responses protocol. The agent container must implement the Foundry Hosted Agent Responses contract (`POST /responses` and `GET /readiness`). The adapter invokes the deployed agent through the -project endpoint; it is not the Hosted Agent container itself. - -The Hosted Agent must honor function tools supplied on each Responses request -for Orka brokered-tool mode. Brokered profiles are disabled by default; enable -only the classes the deployed agent has passed in conformance. If the agent -ignores request-provided function tools, keep observed mode or update the agent -implementation. Do not move Orka -production tool credentials into Foundry Toolbox or MCP merely to make a probe -pass; that changes the governance boundary. +project endpoint. The default adapter runs outside Foundry; the optional +v2 hosted package runs its supervisor inside a separate Hosted Agent. + +For Orka brokered-tool mode, use one of two schema delivery modes. The default +`request` mode requires the Hosted Agent endpoint to accept function tools on +each Responses request. `provider-static` mode is available for Hosted Agent +endpoints that reject request-level tools; in that mode the Hosted Agent must +preconfigure the function schemas. The adapter still rejects any function call +that was not supplied in the current Orka turn, and Orka still owns argument +validation, policy, approvals, credentials, execution, and audit. Brokered +profiles are disabled by default; enable only the classes and schema mode the +deployed agent has passed in conformance. Do not move Orka production tool +credentials into Foundry Toolbox or MCP merely to make a probe pass; that +changes the governance boundary. ## Configuration @@ -56,7 +66,8 @@ pass; that changes the governance boundary. | `ORKA_FOUNDRY_TURN_TIMEOUT` | Absolute adapter maximum for one Orka turn, default `20s`. | | `ORKA_FOUNDRY_ISOLATION_MODE` | `entra` (default) or `header`. In `header` mode, the adapter sends an opaque hash of Orka's runtime session ID as `x-ms-user-isolation-key`. | | `ORKA_FOUNDRY_FEATURES` | Preview feature header value, default `HostedAgents=V1Preview`. Set an empty value only when the deployed API no longer requires the header. | -| `ORKA_FOUNDRY_BROKERED_TOOL_CLASSES` | Optional comma-separated classes to advertise and accept: `read`, `write`, or `read,write`. Empty by default (observed-only). Enable only after the Hosted Agent passes the matching conformance probes with request-provided function tools. | +| `ORKA_FOUNDRY_BROKERED_TOOL_CLASSES` | Optional comma-separated classes to advertise and accept: `read`, `write`, or `read,write`. Empty by default (observed-only). Enable only after the Hosted Agent passes the matching conformance probes in the configured schema mode. | +| `ORKA_FOUNDRY_TOOL_SCHEMA_MODE` | Brokered schema delivery: `request` (default) sends Orka's safe function schemas on each Responses request; `provider-static` omits request-level tools and requires matching schemas to be preconfigured in the Hosted Agent. | The adapter authenticates with Azure SDK `DefaultAzureCredential` and requests the `https://ai.azure.com/.default` scope. In Kubernetes, use Azure Workload @@ -108,7 +119,8 @@ spec: `store: true`. - `response.created` becomes `TurnStarted`. - `response.output_text.delta` becomes `RuntimeOutput`. -- When explicitly enabled, request-provided safe Orka tool schemas become Responses function tools. +- When explicitly enabled in `request` mode, request-provided safe Orka tool schemas become Responses function tools. +- In `provider-static` mode, the provider uses its preconfigured schemas while the adapter enforces every returned function name against the current turn's Orka-supplied allowlist. - `function_call` output items become `ToolCallRequested` frames. - `/v1/turns/{turnID}/continue` sends `function_call_output` items with the same `call_id` and chains them with `previous_response_id`. diff --git a/acp.go b/acp.go new file mode 100644 index 0000000..c899a41 --- /dev/null +++ b/acp.go @@ -0,0 +1,367 @@ +package main + +import ( + "bufio" + "context" + "crypto/rand" + "encoding/hex" + "encoding/json" + "errors" + "io" + "net/http" + "os" + "path/filepath" + "sync" + "unicode/utf8" +) + +type acpSession struct { + id string + mcp *acpMCPClient + previous string + poisoned bool +} + +type acpOperation struct { + key string + id json.RawMessage + ctx context.Context + cancel context.CancelFunc +} + +type acpWrite struct { + frame []byte + done chan error +} + +type acpServer struct { + ctx context.Context + cancel context.CancelFunc + cfg acpConfiguration + client *http.Client + writes chan acpWrite + wg sync.WaitGroup + + mu sync.Mutex + initialized bool + session *acpSession + active *acpOperation +} + +func serveACP(parent context.Context, cfg acpConfiguration, input io.ReadCloser, output io.WriteCloser) error { + ctx, cancel := context.WithCancel(parent) + s := &acpServer{ctx: ctx, cancel: cancel, cfg: cfg, client: newACPHTTPClient(), writes: make(chan acpWrite, 32)} + type readResult struct { + line []byte + err error + } + reads := make(chan readResult, 1) + var transport sync.WaitGroup + transport.Add(2) + go func() { + defer transport.Done() + reader := bufio.NewReaderSize(input, 32<<10) + for { + line, err := acpReadLine(reader) + select { + case reads <- readResult{line, err}: + case <-ctx.Done(): + return + } + if err != nil { + return + } + } + }() + go func() { + defer transport.Done() + for { + select { + case <-ctx.Done(): + return + case write := <-s.writes: + n, err := output.Write(write.frame) + if err == nil && n != len(write.frame) { + err = io.ErrShortWrite + } + write.done <- err + if err != nil { + cancel() + return + } + } + } + }() + defer func() { + cancel() + _ = input.Close() + _ = output.Close() + s.wg.Wait() + transport.Wait() + s.client.CloseIdleConnections() + }() + for { + select { + case <-ctx.Done(): + return errACPTransport + case read := <-reads: + if errors.Is(read.err, io.EOF) { + return nil + } + if read.err != nil { + return errACPTransport + } + if len(read.line) != 0 { + s.accept(read.line) + } + } + } +} + +func (s *acpServer) enqueue(value any) (<-chan error, error) { + frame, err := json.Marshal(value) + if err != nil || len(frame) >= acpMaxMessageBytes { + return nil, errACPTransport + } + write := acpWrite{frame: append(frame, '\n'), done: make(chan error, 1)} + select { + case s.writes <- write: + return write.done, nil + case <-s.ctx.Done(): + return nil, errACPTransport + default: + // A bounded queue lets cancellation continue to read stdin even when + // the parent stops consuming stdout. Floods close the child channel. + s.cancel() + return nil, errACPTransport + } +} + +func (s *acpServer) send(value any) error { + done, err := s.enqueue(value) + if err != nil { + return err + } + select { + case err := <-done: + return err + case <-s.ctx.Done(): + return errACPTransport + } +} + +func (s *acpServer) respond(id json.RawMessage, result any, failure *acpRPCError) { + if len(id) == 0 { + id = json.RawMessage("null") + } + if _, err := s.enqueue(acpResponse{JSONRPC: "2.0", ID: id, Result: result, Error: failure}); err != nil { + s.cancel() + } +} + +func (s *acpServer) accept(line []byte) { + var request acpRequest + if !json.Valid(line) { + s.respond(nil, nil, &acpRPCError{-32700, "invalid ACP JSON"}) + return + } + if acpDecode(line, &request, true) != nil || request.JSONRPC != "2.0" || !acpSafeString(request.Method, 128) { + s.respond(nil, nil, acpInvalidRequest) + return + } + if len(request.ID) == 0 { + s.notification(request) + return + } + key, err := acpRequestKey(request.ID) + if err != nil { + s.respond(nil, nil, acpInvalidRequest) + return + } + s.mu.Lock() + defer s.mu.Unlock() + if s.active != nil && s.active.key == key { + // Two live requests cannot share one response identity. + s.active.cancel() + s.cancel() + return + } + switch request.Method { + case "initialize": + var params struct { + ProtocolVersion int `json:"protocolVersion"` + ClientCapabilities json.RawMessage `json:"clientCapabilities"` + ClientInfo json.RawMessage `json:"clientInfo"` + Meta json.RawMessage `json:"_meta,omitempty"` + } + if s.initialized || acpDecode(request.Params, ¶ms, true) != nil || params.ProtocolVersion != 1 || + (len(params.ClientCapabilities) != 0 && params.ClientCapabilities[0] != '{') { + s.respond(request.ID, nil, acpInvalidParams) + return + } + s.initialized = true + s.respond(request.ID, map[string]any{ + "protocolVersion": 1, + "agentCapabilities": map[string]any{ + "loadSession": false, + "promptCapabilities": map[string]bool{"image": false, "audio": false, "embeddedContext": false}, + "mcpCapabilities": map[string]bool{"http": true}, + "sessionCapabilities": map[string]any{}, "auth": map[string]any{}, + }, + "agentInfo": map[string]string{"name": "foundry-acp", "version": "1"}, + }, nil) + case "session/new": + s.newSessionLocked(request, key) + case "session/prompt": + s.promptLocked(request, key) + default: + s.respond(request.ID, nil, &acpRPCError{-32601, "ACP method not supported"}) + } +} + +func (s *acpServer) notification(request acpRequest) { + s.mu.Lock() + defer s.mu.Unlock() + if s.active == nil { + return + } + switch request.Method { + case "session/cancel": + var params struct { + SessionID string `json:"sessionId"` + Meta json.RawMessage `json:"_meta,omitempty"` + } + if acpDecode(request.Params, ¶ms, true) == nil && s.session != nil && params.SessionID == s.session.id { + s.active.cancel() + } + case "$/cancel_request": + var params struct { + RequestID json.RawMessage `json:"requestId"` + } + key, err := "", error(nil) + if acpDecode(request.Params, ¶ms, true) == nil { + key, err = acpRequestKey(params.RequestID) + } + if err == nil && key == s.active.key { + s.active.cancel() + } + } +} + +func (s *acpServer) newSessionLocked(request acpRequest, key string) { + if !s.initialized || s.session != nil || s.active != nil { + s.respond(request.ID, nil, acpInvalidRequest) + return + } + var params acpNewSession + if acpDecode(request.Params, ¶ms, true) != nil || len(params.AdditionalDirectories) != 0 || len(params.MCPServers) != 1 || !filepath.IsAbs(params.CWD) { + s.respond(request.ID, nil, acpInvalidParams) + return + } + cwd, err := os.Getwd() + requested, requestedErr := filepath.EvalSymlinks(params.CWD) + actual, actualErr := filepath.EvalSymlinks(cwd) + if err != nil || requestedErr != nil || actualErr != nil || requested != actual { + s.respond(request.ID, nil, acpInvalidParams) + return + } + mcp, err := newACPMCPClient(params.MCPServers[0], s.client) + if err != nil { + s.respond(request.ID, nil, acpInvalidParams) + return + } + op := s.newOperationLocked(request, key) + s.wg.Go(func() { + defer op.cancel() + err := mcp.initialize(op.ctx) + s.mu.Lock() + defer s.mu.Unlock() + if err != nil || op.ctx.Err() != nil { + s.respond(op.id, nil, acpInternalError) + } else { + s.session = &acpSession{id: acpOpaqueID("foundry-"), mcp: mcp} + s.respond(op.id, map[string]string{"sessionId": s.session.id}, nil) + } + s.active = nil + }) +} + +func (s *acpServer) newOperationLocked(request acpRequest, key string) *acpOperation { + ctx, cancel := context.WithCancel(s.ctx) + op := &acpOperation{key: key, id: request.ID, ctx: ctx, cancel: cancel} + s.active = op + return op +} + +func (s *acpServer) promptLocked(request acpRequest, key string) { + var params acpPrompt + if acpDecode(request.Params, ¶ms, true) != nil { + s.respond(request.ID, nil, acpInvalidParams) + return + } + text, err := params.text() + if err != nil || s.session == nil || params.SessionID != s.session.id || s.session.poisoned { + s.respond(request.ID, nil, acpInvalidParams) + return + } + if s.active != nil { + s.respond(request.ID, nil, &acpRPCError{-32800, "ACP session already has an active prompt"}) + return + } + op := s.newOperationLocked(request, key) + session := s.session + s.wg.Go(func() { + defer op.cancel() + previous, output, err := s.runPrompt(op.ctx, session, text) + if err == nil { + err = s.outputText(op.ctx, session.id, output) + } + s.mu.Lock() + defer s.mu.Unlock() + switch { + case op.ctx.Err() != nil: + session.poisoned, session.previous = true, "" + s.respond(op.id, map[string]string{"stopReason": "cancelled"}, nil) + case err != nil: + session.poisoned, session.previous = true, "" + s.respond(op.id, nil, acpInternalError) + default: + session.previous = previous + s.respond(op.id, map[string]string{"stopReason": "end_turn"}, nil) + } + s.active = nil + }) +} + +func (s *acpServer) update(sessionID string, update any) error { + return s.send(map[string]any{ + "jsonrpc": "2.0", "method": "session/update", + "params": map[string]any{"sessionId": sessionID, "update": update}, + }) +} + +func (s *acpServer) outputText(ctx context.Context, sessionID, text string) error { + for len(text) != 0 { + if err := ctx.Err(); err != nil { + return err + } + n := min(len(text), acpTextChunkBytes) + for n < len(text) && !utf8.RuneStart(text[n]) { + n-- + } + if err := s.update(sessionID, map[string]any{ + "sessionUpdate": "agent_message_chunk", + "content": map[string]string{"type": "text", "text": text[:n]}, + }); err != nil { + return err + } + text = text[n:] + } + return ctx.Err() +} + +func acpOpaqueID(prefix string) string { + var data [16]byte + _, _ = rand.Read(data[:]) + return prefix + hex.EncodeToString(data[:]) +} diff --git a/acp_config.go b/acp_config.go new file mode 100644 index 0000000..bba21cf --- /dev/null +++ b/acp_config.go @@ -0,0 +1,196 @@ +package main + +import ( + "context" + "crypto/sha256" + "crypto/subtle" + "encoding/hex" + "errors" + "flag" + "io" + "net" + "net/http" + "net/url" + "os" + "strconv" + "strings" + "time" + "unicode" +) + +const ( + acpConfigPath = "/agent/foundry.json" + acpProviderBaseEnv = "ORKA_FOUNDRY_ACP_PROVIDER_BASE_URL" + acpProviderTokenEnv = "ORKA_FOUNDRY_ACP_PROVIDER_TOKEN" + acpModelEnv = "ORKA_FOUNDRY_ACP_MODEL" + acpConfigDigestEnv = "ORKA_FOUNDRY_ACP_AGENT_CONFIGURATION_DIGEST" + acpMaxConfigBytes = 64 << 10 + acpMaxMessageBytes = 8 << 20 + acpTextChunkBytes = 32 << 10 + acpHTTPTimeout = 120 * time.Second +) + +var ( + errACPConfiguration = errors.New("invalid Foundry ACP configuration") + errACPProvider = errors.New("Foundry ACP provider request failed") + errACPMCP = errors.New("Foundry ACP MCP request failed") + errACPTransport = errors.New("Foundry ACP transport failed") +) + +type acpAgentConfiguration struct { + Model string `json:"model"` + ToolSchemaMode string `json:"toolSchemaMode"` + HostedTarget acpHostedTarget `json:"hostedTarget"` +} + +type acpHostedTarget struct { + ProjectEndpoint string `json:"projectEndpoint"` + AgentName string `json:"agentName"` + AgentVersion string `json:"agentVersion"` +} + +type acpConfiguration struct { + agent acpAgentConfiguration + providerURL string + token string +} + +// The ACP entry point deliberately precedes Azure credential initialization. +// Its only network authority is the two supervisor-owned loopback proxies. +func maybeServeACP(args []string, input io.ReadCloser, output io.WriteCloser) (bool, error) { + selected := false + for _, arg := range args { + if arg == "--protocol" || strings.HasPrefix(arg, "--protocol=") { + selected = true + break + } + } + if !selected { + return false, nil + } + flags := flag.NewFlagSet("foundry-acp", flag.ContinueOnError) + flags.SetOutput(io.Discard) + protocol := flags.String("protocol", "", "") + path := flags.String("config", acpConfigPath, "") + if flags.Parse(args) != nil || flags.NArg() != 0 || *protocol != "acp" { + return true, errACPConfiguration + } + cfg, err := loadACPConfiguration(*path, os.Getenv) + if err != nil { + return true, err + } + return true, serveACP(context.Background(), cfg, input, output) +} + +func loadACPConfiguration(path string, getenv func(string) string) (acpConfiguration, error) { + file, err := os.Open(path) + if err != nil { + return acpConfiguration{}, errACPConfiguration + } + defer file.Close() //nolint:errcheck + data, err := io.ReadAll(io.LimitReader(file, acpMaxConfigBytes+1)) + if err != nil || len(data) > acpMaxConfigBytes { + return acpConfiguration{}, errACPConfiguration + } + return verifyACPConfiguration(data, getenv) +} + +func verifyACPConfiguration(data []byte, getenv func(string) string) (acpConfiguration, error) { + agent, err := decodeACPAgentConfiguration(data, getenv(acpConfigDigestEnv), getenv(acpModelEnv)) + if err != nil { + return acpConfiguration{}, err + } + base, err := acpLoopbackURL(getenv(acpProviderBaseEnv)) + token := getenv(acpProviderTokenEnv) + if err != nil || !acpSafeString(token, 16<<10) || strings.ContainsAny(token, " \t") { + return acpConfiguration{}, errACPConfiguration + } + base.Path = strings.TrimRight(base.Path, "/") + "/responses" + return acpConfiguration{agent: agent, providerURL: base.String(), token: token}, nil +} + +// Both entry points verify one immutable buffer, while only the privileged +// broker uses HostedTarget. No child proxy credential is needed to parse it. +func decodeACPAgentConfiguration(data []byte, expectedDigest, expectedModel string) (acpAgentConfiguration, error) { + actual := sha256.Sum256(data) + encoded := "sha256:" + hex.EncodeToString(actual[:]) + if len(data) > acpMaxConfigBytes || subtle.ConstantTimeCompare([]byte(expectedDigest), []byte(encoded)) != 1 { + return acpAgentConfiguration{}, errACPConfiguration + } + var agent acpAgentConfiguration + if acpDecode(data, &agent, true) != nil || !acpSafeString(agent.Model, 512) || agent.Model != expectedModel { + return acpAgentConfiguration{}, errACPConfiguration + } + if agent.ToolSchemaMode != toolSchemaModeRequest && agent.ToolSchemaMode != toolSchemaModeProviderStatic { + return acpAgentConfiguration{}, errACPConfiguration + } + if strings.TrimSpace(agent.HostedTarget.ProjectEndpoint) != agent.HostedTarget.ProjectEndpoint || + !foundryEndpointIsSafe(agent.HostedTarget.ProjectEndpoint) || + validateAgentName(agent.HostedTarget.AgentName) != nil || agent.HostedTarget.AgentVersion == "" || + strings.EqualFold(agent.HostedTarget.AgentVersion, "latest") || validateAgentVersion(agent.HostedTarget.AgentVersion) != nil { + return acpAgentConfiguration{}, errACPConfiguration + } + return agent, nil +} + +func acpSafeString(value string, limit int) bool { + if value == "" || len(value) > limit { + return false + } + for _, ch := range value { + if unicode.IsControl(ch) { + return false + } + } + return true +} + +func acpLoopbackURL(value string) (*url.URL, error) { + u, err := url.Parse(value) + if err != nil || !acpSafeString(value, 8<<10) || strings.TrimSpace(value) != value || u == nil || + (u.Scheme != "http" && u.Scheme != "https") || u.Host == "" || u.User != nil || + u.RawQuery != "" || u.ForceQuery || u.Fragment != "" || strings.Contains(value, "#") || u.RawPath != "" { + return nil, errACPConfiguration + } + ip := net.ParseIP(u.Hostname()) + if u.Hostname() != "localhost" && (ip == nil || !ip.IsLoopback()) { + return nil, errACPConfiguration + } + if port := u.Port(); port != "" { + value, err := strconv.Atoi(port) + if err != nil || value < 1 || value > 65535 { + return nil, errACPConfiguration + } + } + return u, nil +} + +func newACPHTTPClient() *http.Client { + dialer := &net.Dialer{Timeout: 10 * time.Second, KeepAlive: 30 * time.Second} + return &http.Client{ + Timeout: acpHTTPTimeout, + CheckRedirect: func(*http.Request, []*http.Request) error { + return errACPTransport + }, + Transport: &http.Transport{ + Proxy: nil, + DialContext: func(ctx context.Context, network, address string) (net.Conn, error) { + host, port, err := net.SplitHostPort(address) + if err != nil { + return nil, errACPTransport + } + if host == "localhost" { + host = "127.0.0.1" + } + if ip := net.ParseIP(host); ip == nil || !ip.IsLoopback() { + return nil, errACPTransport + } + return dialer.DialContext(ctx, network, net.JoinHostPort(host, port)) + }, + MaxIdleConns: 4, + MaxIdleConnsPerHost: 2, + MaxConnsPerHost: 2, + IdleConnTimeout: 30 * time.Second, + }, + } +} diff --git a/acp_config_test.go b/acp_config_test.go new file mode 100644 index 0000000..034f5e5 --- /dev/null +++ b/acp_config_test.go @@ -0,0 +1,189 @@ +package main + +import ( + "crypto/sha256" + "encoding/hex" + "encoding/json" + "io" + "net/http" + "net/http/httptest" + "strings" + "sync/atomic" + "testing" +) + +func acpTestConfigBytes(mode string) []byte { + data, _ := json.Marshal(acpAgentConfiguration{ + Model: "test-model", ToolSchemaMode: mode, + HostedTarget: acpHostedTarget{ProjectEndpoint: "https://foundry.example/api/projects/test", AgentName: "test-agent", AgentVersion: "7"}, + }) + return append(data, '\n') +} + +func acpTestDigest(data []byte) string { + digest := sha256.Sum256(data) + return "sha256:" + hex.EncodeToString(digest[:]) +} + +func acpTestEnvironment(data []byte) map[string]string { + return map[string]string{ + acpConfigDigestEnv: acpTestDigest(data), acpModelEnv: "test-model", + acpProviderBaseEnv: "http://127.0.0.1:1234/local/v1", acpProviderTokenEnv: "test-only-proxy-token", + } +} + +func TestACPConfigurationPinsExactBytesAndHostedTarget(t *testing.T) { + data := acpTestConfigBytes(toolSchemaModeProviderStatic) + env := acpTestEnvironment(data) + cfg, err := verifyACPConfiguration(data, func(key string) string { return env[key] }) + if err != nil || cfg.providerURL != "http://127.0.0.1:1234/local/v1/responses" || cfg.agent.HostedTarget.AgentVersion != "7" { + t.Fatal("valid pinned configuration rejected") + } + for name, mutate := range map[string]func(map[string]string){ + "wrong digest": func(e map[string]string) { e[acpConfigDigestEnv] = "sha256:" + strings.Repeat("0", 64) }, + "uppercase digest": func(e map[string]string) { e[acpConfigDigestEnv] = strings.ToUpper(e[acpConfigDigestEnv]) }, + "wrong model": func(e map[string]string) { e[acpModelEnv] = "other" }, + "missing token": func(e map[string]string) { delete(e, acpProviderTokenEnv) }, + "newline token": func(e map[string]string) { e[acpProviderTokenEnv] = "test\nvalue" }, + "public provider": func(e map[string]string) { e[acpProviderBaseEnv] = "https://foundry.example/v1" }, + "query provider": func(e map[string]string) { e[acpProviderBaseEnv] += "?credential=test" }, + "fragment provider": func(e map[string]string) { e[acpProviderBaseEnv] += "#" }, + "userinfo provider": func(e map[string]string) { e[acpProviderBaseEnv] = "http://user@127.0.0.1/v1" }, + "invalid provider port": func(e map[string]string) { e[acpProviderBaseEnv] = "http://127.0.0.1:65536/v1" }, + } { + t.Run(name, func(t *testing.T) { + e := acpTestEnvironment(data) + mutate(e) + if _, err := verifyACPConfiguration(data, func(key string) string { return e[key] }); err == nil { + t.Fatal("unsafe configuration accepted") + } + }) + } + if _, err := verifyACPConfiguration(append(data, '\n'), func(key string) string { return env[key] }); err == nil { + t.Fatal("different raw configuration bytes accepted") + } +} + +func TestACPMCPServerRejectsUnauthenticatedAndNonlocalTargets(t *testing.T) { + valid := `{"type":"http","name":"broker","url":"http://127.0.0.1:1234/mcp","headers":[{"name":"Authorization","value":"Bearer test-only"}]}` + for name, data := range map[string]string{ + "nonlocal": strings.Replace(valid, "127.0.0.1", "10.0.0.1", 1), + "remote hostname": strings.Replace(valid, "127.0.0.1", "remote.invalid", 1), + "no header": strings.Replace(valid, `[{"name":"Authorization","value":"Bearer test-only"}]`, `[]`, 1), + "empty bearer": strings.Replace(valid, "Bearer test-only", "Bearer ", 1), + "invalid bearer": strings.Replace(valid, "Bearer test-only", "Basic test-only", 1), + "SSE transport": strings.Replace(valid, `"http"`, `"sse"`, 1), + "process command": strings.Replace(valid, `"type":`, `"command":"forbidden","type":`, 1), + "extra header": strings.Replace(valid, `"headers":[`, `"headers":[{"name":"X-Override","value":"test"},`, 1), + "duplicate authorization": strings.Replace(valid, `"headers":[`, `"headers":[{"name":"authorization","value":"Bearer other"},`, 1), + } { + t.Run(name, func(t *testing.T) { + var server acpMCPServer + if err := acpDecode([]byte(data), &server, true); err != nil { + return + } + if _, err := newACPMCPClient(server, newACPHTTPClient()); err == nil { + t.Fatal("unsafe MCP endpoint accepted") + } + }) + } +} + +func TestACPConfigurationRejectsBakedToolsAndUnpinnedTargets(t *testing.T) { + base := string(acpTestConfigBytes(toolSchemaModeRequest)) + cases := map[string]string{ + "tools": strings.Replace(base, `"model":`, `"tools":[],"model":`, 1), + "brokeredTools": strings.Replace(base, `"model":`, `"brokeredTools":[],"model":`, 1), + "context": strings.Replace(base, `"model":`, `"context":{},"model":`, 1), + "unknown mode": strings.Replace(base, `"request"`, `"native"`, 1), + "missing target": `{"model":"test-model","toolSchemaMode":"request"}`, + "missing version": strings.Replace(base, `"agentVersion":"7"`, `"agentVersion":""`, 1), + "floating version": strings.Replace(base, `"agentVersion":"7"`, `"agentVersion":"latest"`, 1), + "alias version": strings.Replace(base, `"agentVersion":"7"`, `"agentVersion":"@latest"`, 1), + "unknown target field": strings.Replace(base, `"agentVersion":"7"`, `"agentVersion":"7","sessionId":"untrusted"`, 1), + "duplicate key": strings.Replace(base, `"model":`, `"model":"other","model":`, 1), + } + for name, raw := range cases { + t.Run(name, func(t *testing.T) { + data := []byte(raw) + if _, err := decodeACPAgentConfiguration(data, acpTestDigest(data), "test-model"); err == nil { + t.Fatal("unsupported configuration accepted") + } + }) + } +} + +func TestACPHTTPTransportRejectsRedirectsAndEnvironmentProxy(t *testing.T) { + var redirected, proxied atomic.Int32 + destination := httptest.NewServer(http.HandlerFunc(func(w http.ResponseWriter, _ *http.Request) { + redirected.Add(1) + w.WriteHeader(http.StatusOK) + })) + defer destination.Close() + proxy := httptest.NewServer(http.HandlerFunc(func(w http.ResponseWriter, _ *http.Request) { + proxied.Add(1) + w.WriteHeader(http.StatusOK) + })) + defer proxy.Close() + t.Setenv("HTTP_PROXY", proxy.URL) + t.Setenv("HTTPS_PROXY", proxy.URL) + source := httptest.NewServer(http.HandlerFunc(func(w http.ResponseWriter, r *http.Request) { + http.Redirect(w, r, destination.URL, http.StatusTemporaryRedirect) + })) + defer source.Close() + client := newACPHTTPClient() + defer client.CloseIdleConnections() + response, err := client.Get(source.URL) + if response != nil { + _ = response.Body.Close() + } + if err == nil || redirected.Load() != 0 || proxied.Load() != 0 { + t.Fatal("local transport followed redirect or consulted environment proxy") + } + response, err = client.Get("http://not-loopback.invalid/") + if response != nil { + _ = response.Body.Close() + } + if err == nil || proxied.Load() != 0 { + t.Fatal("local transport attempted non-loopback access") + } +} + +func TestACPStrictJSONAndFraming(t *testing.T) { + for _, data := range []string{ + `{"a":1,"a":2}`, `{"a":{"b":1,"b":2}}`, `{"a":1} {"b":2}`, + strings.Repeat("[", 66) + "0" + strings.Repeat("]", 66), string([]byte{'"', 0xff, '"'}), + `"\ud800"`, `"\udfff"`, `"\ud800\u0000"`, + } { + var value any + if acpDecode([]byte(data), &value, false) == nil { + t.Fatal("ambiguous or unbounded JSON accepted") + } + } + for _, raw := range []string{`"\ud83c\udf0d"`, `"\\ud800"`, `"\ufffd"`, `"héllo 🌍"`} { + var value string + if acpDecode([]byte(raw), &value, false) != nil { + t.Fatal("valid Unicode rejected") + } + } + for _, raw := range []string{`null`, `true`, `1.5`, `1e0`, `[]`, `{}`, `""`} { + if _, err := acpRequestKey(json.RawMessage(raw)); err == nil { + t.Fatal("unsupported request identity accepted") + } + } + for _, raw := range []string{`1`, `-2`, `"request-1"`} { + if _, err := acpRequestKey(json.RawMessage(raw)); err != nil { + t.Fatal("supported request identity rejected") + } + } +} + +func TestACPProtocolSelectionDoesNotInitializeAzure(t *testing.T) { + in := io.NopCloser(strings.NewReader("")) + if handled, err := maybeServeACP(nil, in, nil); handled || err != nil { + t.Fatal("legacy entry point was claimed") + } + if handled, err := maybeServeACP([]string{"--protocol", "unsupported"}, in, nil); !handled || err == nil { + t.Fatal("unsupported protocol was not rejected") + } +} diff --git a/acp_folded_response_test.go b/acp_folded_response_test.go new file mode 100644 index 0000000..c801685 --- /dev/null +++ b/acp_folded_response_test.go @@ -0,0 +1,72 @@ +package main + +import ( + "encoding/json" + "fmt" + "net/http" + "strings" + "sync/atomic" + "testing" +) + +func TestACPFoldedResponseFieldsCannotExecuteTool(t *testing.T) { + const call = `{"id":"item-1","type":"function_call","name":"probe","call_id":"call-1","arguments":"{}","status":"completed"}` + for name, document := range map[string]string{ + "response status": `{"id":"response-1","status":"in_progress","Status":"completed","output":[` + call + `]}`, + "failed response": `{"id":"response-1","status":"failed","Status":"completed","error":{"code":"fixture-failure"},"Error":null,"output":[` + call + `]}`, + "incomplete response": `{"id":"response-1","status":"completed","incomplete_details":{"reason":"max_output_tokens"},"Incomplete_Details":null,"output":[` + call + `]}`, + "native item type": `{"id":"response-1","status":"completed","output":[` + strings.Replace(call, `"type":"function_call"`, `"type":"web_search_call","Type":"function_call"`, 1) + `]}`, + "incomplete item status": `{"id":"response-1","status":"completed","output":[` + strings.Replace(call, `"status":"completed"`, `"status":"in_progress","Status":"completed"`, 1) + `]}`, + } { + for _, mediaType := range []string{"application/json", "text/event-stream"} { + t.Run(name+"/"+mediaType, func(t *testing.T) { + var requests atomic.Int32 + mcp := &acpTestMCP{ + tools: func() []map[string]any { return acpTestTools("probe") }, + execute: func(w http.ResponseWriter, _ *http.Request, id json.RawMessage, _ string, _ json.RawMessage) { + acpTestToolResult(w, id, "unexpected", false) + }, + } + peer := newACPTestPeer(t, toolSchemaModeProviderStatic, func(w http.ResponseWriter, r *http.Request) { + acpTestReadProvider(t, r) + if requests.Add(1) > 1 { + acpTestCompleted(w, "response-2", "unexpected") + return + } + w.Header().Set("Content-Type", mediaType) + if mediaType == "application/json" { + _, _ = fmt.Fprint(w, document) + } else { + _, _ = fmt.Fprint(w, acpTestSSE(`{"type":"response.completed","response":`+document+`}`)) + } + }, mcp) + reply := peer.reply(peer.prompt("folded semantic fields")) + if requests.Load() != 1 || mcp.calls.Load() != 0 || len(peer.events) != 0 { + t.Fatalf("malformed response admitted effects: providerRequests=%d toolCalls=%d events=%d", requests.Load(), mcp.calls.Load(), len(peer.events)) + } + acpAssertFailure(t, reply) + }) + } + } +} + +func TestACPFoldedResponseNestedFieldsRejected(t *testing.T) { + for name, document := range map[string]string{ + "response identity": `{"id":"wrong","ID":"response-1","status":"completed"}`, + "Session identity": `{"id":"response-1","status":"completed","agent_session_id":"wrong","Agent_Session_ID":"owned"}`, + "item identity": `{"id":"response-1","status":"completed","output":[{"id":"wrong","ID":"item-1","type":"message"}]}`, + "item role": `{"id":"response-1","status":"completed","output":[{"type":"message","role":"user","Role":"assistant","content":[{"type":"output_text","text":"fixture"}]}]}`, + "content type": `{"id":"response-1","status":"completed","output":[{"type":"message","content":[{"type":"refusal","Type":"output_text","text":"fixture"}]}]}`, + "Unicode status": `{"id":"response-1","status":"in_progress","ſtatus":"completed"}`, + "escaped status": `{"id":"response-1","status":"in_progress","\u0053tatus":"completed"}`, + } { + t.Run(name, func(t *testing.T) { + if _, err := acpDecodeFoundryResponse([]byte(document)); err == nil { + t.Error("ambiguous response or nested field accepted") + } + if _, err := acpParseFoundrySSE(strings.NewReader(acpTestSSE(`{"type":"response.completed","response":` + document + `}`))); err == nil { + t.Error("ambiguous SSE response or nested field accepted") + } + }) + } +} diff --git a/acp_json_fields.go b/acp_json_fields.go new file mode 100644 index 0000000..29476a9 --- /dev/null +++ b/acp_json_fields.go @@ -0,0 +1,95 @@ +package main + +import ( + "encoding/json" + "reflect" + "strings" +) + +type acpJSONStructField struct { + name string + typeOf reflect.Type + depth int + tagged bool +} + +func acpJSONStructType(value reflect.Type) reflect.Type { + for value != nil && value.Kind() == reflect.Pointer { + value = value.Elem() + } + if value == nil || reflect.PointerTo(value).Implements(reflect.TypeFor[json.Unmarshaler]()) { + return nil + } + return value +} + +// Only struct fields define folded names. In particular, RawMessage tool +// arguments, map keys and unknown metadata retain their case-sensitive JSON. +func acpJSONStructFields(value reflect.Type, depth int) ([]acpJSONStructField, error) { + if depth > 64 { + return nil, acpInvalidParams + } + value = acpJSONStructType(value) + if value == nil || value.Kind() != reflect.Struct { + return nil, nil + } + var candidates []acpJSONStructField + for i := 0; i < value.NumField(); i++ { + field := value.Field(i) + name, _, _ := strings.Cut(field.Tag.Get("json"), ",") + if name == "-" || (!field.IsExported() && !field.Anonymous) { + continue + } + if name == "" && field.Anonymous { + if embedded := acpJSONStructType(field.Type); embedded != nil && embedded.Kind() == reflect.Struct { + fields, err := acpJSONStructFields(field.Type, depth+1) + if err != nil { + return nil, err + } + candidates = append(candidates, fields...) + continue + } + } + if !field.IsExported() { + continue + } + tagged := name != "" + if name == "" { + name = field.Name + } + candidates = append(candidates, acpJSONStructField{name, field.Type, depth, tagged}) + } + // Match encoding/json's dominance rules for the selected plain structs: + // shallower fields win; at equal depth a tagged field wins; ties are ignored. + var fields []acpJSONStructField + for i, candidate := range candidates { + keep := true + for j, other := range candidates { + if i == j || candidate.name != other.name { + continue + } + if other.depth < candidate.depth || (other.depth == candidate.depth && (other.tagged || !candidate.tagged)) { + keep = false + break + } + } + if keep { + fields = append(fields, candidate) + } + } + return fields, nil +} + +func acpJSONMatchField(fields []acpJSONStructField, name string) (string, reflect.Type) { + for _, field := range fields { + if field.name == name { + return field.name, field.typeOf + } + } + for _, field := range fields { + if strings.EqualFold(field.name, name) { + return field.name, field.typeOf + } + } + return name, nil +} diff --git a/acp_json_fields_test.go b/acp_json_fields_test.go new file mode 100644 index 0000000..0c7b429 --- /dev/null +++ b/acp_json_fields_test.go @@ -0,0 +1,120 @@ +package main + +import ( + "bytes" + "encoding/json" + "strings" + "testing" +) + +func TestACPStructFieldsKeepOpaqueValuesCaseSensitive(t *testing.T) { + var value struct { + foundryOutputItem + Status string `json:"status"` + Map map[string]any `json:"map"` + Any any `json:"any"` + } + data := []byte(`{"TyPe":"function_call","ſtatus":"completed","ArGuMeNtS":{"Key":1,"key":2},"Map":{"Key":3,"key":4},"Any":{"Key":5,"key":6},"unknown":{"Key":7,"key":8},"Unknown":null}`) + if err := acpDecodeStruct(data, &value, false); err != nil { + t.Fatal("single aliases or opaque case-sensitive values rejected") + } + if value.Type != "function_call" || value.Status != "completed" || + string(value.Arguments) != `{"Key":1,"key":2}` || len(value.Map) != 2 || len(value.Any.(map[string]any)) != 2 { + t.Fatal("structured decode changed an opaque value") + } + for _, data := range []string{ + `{"arguments":{"Key":1,"Key":2}}`, + `{"Map":{"Key":1,"\u004bey":2}}`, + `{"Any":{"nested":{"same":1,"same":2}}}`, + `{"unknown":{"same":1,"same":2}}`, + `{"arguments":"\ud800"}`, + } { + if acpDecodeStruct([]byte(data), &value, false) == nil { + t.Error("opaque value bypassed existing exact duplicate or Unicode validation") + } + } +} + +func TestACPStructFieldsRejectAliasesBeforeApplyingValues(t *testing.T) { + type nested struct { + Status string `json:"status"` + } + for _, data := range []string{ + `{"status":"active","Status":"idle"}`, + `{"status":"active","ſtatus":"idle"}`, + `{"status":"active","\u0053tatus":"idle"}`, + `{"status":null,"Status":"idle"}`, + `{"values":[{"status":"active","Status":"idle"}]}`, + `{"pointer":{"status":"active","Status":"idle"}}`, + `{"pointer":{"status":"active"},"Pointer":{"status":"idle"}}`, + } { + var value struct { + nested + Values []nested `json:"values"` + Pointer *nested `json:"pointer"` + } + value.Status = "untouched" + if acpDecodeStruct([]byte(data), &value, false) == nil || value.Status != "untouched" || value.Pointer != nil || value.Values != nil { + t.Error("folded alias was accepted or applied before rejection") + } + } +} + +func TestACPStructFieldsPreserveDecoderCompatibility(t *testing.T) { + type embedded struct { + ID string `json:"id"` + } + for _, data := range []string{ + `{"ID":"first","id":"second","Status":"completed"}`, + `{"Id":"folded","status":"completed"}`, + `{"id":"canonical","unknown":{"Status":"active","status":"idle"}}`, + } { + var ordinary, checked struct { + embedded + ID string `json:"ID"` + Status string `json:"status"` + } + if acpDecode([]byte(data), &ordinary, false) != nil || acpDecodeStruct([]byte(data), &checked, false) != nil || ordinary != checked { + t.Error("unambiguous fields differ from encoding/json") + } + } + var object struct { + Value json.RawMessage `json:"value"` + } + if acpDecodeStruct([]byte(`{"unknown":true}`), &object, false) != nil || acpDecodeStruct([]byte(`{"unknown":true}`), &object, true) == nil { + t.Error("unknown-field policy changed") + } + if acpDecodeStruct([]byte(`{"value":`+strings.Repeat("[", 65)+strings.Repeat("]", 65)+`}`), &object, false) == nil { + t.Error("structured decoder lost the depth bound") + } +} + +func TestACPStructFieldsSingleAliasesPreserveResponseAndArguments(t *testing.T) { + data := []byte(`{"ID":"response-1","ſtatus":"completed","OuTpUt":[{"ID":"item-1","TyPe":"function_call","ſtatus":"completed","Name":"probe","CALL_ID":"call-1","Arguments":{"Key":1,"key":2}}]}`) + response, err := acpDecodeFoundryResponse(data) + if err != nil || response.Status != "completed" || len(response.Output) != 1 || + response.Output[0].Type != "function_call" || !bytes.Equal(response.Output[0].Arguments, []byte(`{"Key":1,"key":2}`)) { + t.Fatal("single response/item aliases changed case-sensitive arguments") + } +} + +func TestBrokerResponseEvidenceKeepsUnusableOutputOwnership(t *testing.T) { + const data = `{"id":"response-1","status":"completed","agent_session_id":"owned","output":[{"type":"web_search_call","Type":"function_call"}]}` + response, err := brokerDecodeResponseEvidence([]byte(data)) + if err != nil || response.ID != "response-1" || response.AgentSessionID != "owned" { + t.Fatal("coherent response ownership was discarded because output is unusable") + } + if _, err := acpDecodeFoundryResponse([]byte(data)); err == nil { + t.Fatal("ownership evidence admitted unusable output") + } + for _, data := range []string{ + `{"id":"wrong","ID":"response-1","status":"completed"}`, + `{"id":"response-1","status":"active","Status":"completed"}`, + `{"id":"response-1","status":"completed","agent_session_id":"wrong","Agent_Session_ID":"owned"}`, + `{"id":"response-1","status":"completed","error":{"code":"failure"},"Error":null}`, + } { + if _, err := brokerDecodeResponseEvidence([]byte(data)); err == nil { + t.Error("contradictory fields fabricated response ownership") + } + } +} diff --git a/acp_mcp.go b/acp_mcp.go new file mode 100644 index 0000000..1f97558 --- /dev/null +++ b/acp_mcp.go @@ -0,0 +1,190 @@ +package main + +import ( + "bytes" + "context" + "encoding/json" + "io" + "mime" + "net/http" + "strings" + "sync/atomic" +) + +const acpMCPVersion = "2025-06-18" + +// acpMCPClient speaks only to Orka's per-session loopback proxy. Its route and +// scoped bearer select the session; replies are JSON, without MCP-Session-Id +// negotiation or an SSE stream. +type acpMCPClient struct { + url string + bearer string + client *http.Client + nextID atomic.Uint64 +} + +func newACPMCPClient(server acpMCPServer, client *http.Client) (*acpMCPClient, error) { + if server.Type != "http" || !acpSafeString(server.Name, 128) || len(server.Headers) != 1 { + return nil, acpInvalidParams + } + if _, err := acpLoopbackURL(server.URL); err != nil { + return nil, acpInvalidParams + } + header := server.Headers[0] + value, ok := strings.CutPrefix(header.Value, "Bearer ") + if !strings.EqualFold(header.Name, "Authorization") || !ok || + !acpSafeString(value, 16<<10) || strings.ContainsAny(value, " \t") { + return nil, acpInvalidParams + } + return &acpMCPClient{url: server.URL, bearer: header.Value, client: client}, nil +} + +func (m *acpMCPClient) initialize(ctx context.Context) error { + result, err := m.call(ctx, "initialize", map[string]any{ + "protocolVersion": acpMCPVersion, + "capabilities": map[string]any{}, + "clientInfo": map[string]string{"name": "foundry-acp", "version": "1"}, + }) + if err != nil { + return err + } + var reply struct { + ProtocolVersion string `json:"protocolVersion"` + Capabilities struct { + Tools json.RawMessage `json:"tools"` + } `json:"capabilities"` + } + if acpDecode(result, &reply, false) != nil || reply.ProtocolVersion != acpMCPVersion || + len(reply.Capabilities.Tools) == 0 || reply.Capabilities.Tools[0] != '{' { + return errACPMCP + } + body, _ := json.Marshal(acpRequest{JSONRPC: "2.0", Method: "notifications/initialized"}) + response, err := m.post(ctx, body) + if err != nil { + return err + } + defer response.Body.Close() //nolint:errcheck + if response.StatusCode != http.StatusAccepted && response.StatusCode != http.StatusNoContent { + return errACPMCP + } + return nil +} + +func (m *acpMCPClient) tools(ctx context.Context) ([]foundryToolSchema, error) { + result, err := m.call(ctx, "tools/list", map[string]any{}) + if err != nil { + return nil, err + } + var reply struct { + Tools []struct { + Name string `json:"name"` + Description string `json:"description"` + InputSchema json.RawMessage `json:"inputSchema"` + } `json:"tools"` + NextCursor json.RawMessage `json:"nextCursor"` + } + if acpDecode(result, &reply, false) != nil || reply.Tools == nil || len(reply.Tools) > defaultMaxBrokeredCalls || + (len(reply.NextCursor) != 0 && !bytes.Equal(reply.NextCursor, []byte("null"))) { + return nil, errACPMCP + } + tools := make([]foundryToolSchema, 0, len(reply.Tools)) + seen := make(map[string]bool) + for _, tool := range reply.Tools { + if validateFoundryFunctionName(tool.Name) != nil || seen[tool.Name] || len(tool.InputSchema) > maxFoundryToolSchemaBytes { + return nil, errACPMCP + } + var schema map[string]any + if acpDecode(tool.InputSchema, &schema, false) != nil || schema == nil || schema["type"] != "object" { + return nil, errACPMCP + } + seen[tool.Name] = true + tools = append(tools, foundryToolSchema{Type: "function", Name: tool.Name, Description: tool.Description, Parameters: tool.InputSchema}) + } + encoded, err := json.Marshal(tools) + if err != nil || len(encoded) > maxFoundryToolSchemaBytes { + return nil, errACPMCP + } + return tools, nil +} + +func (m *acpMCPClient) execute(ctx context.Context, name string, args json.RawMessage) (string, bool, error) { + result, err := m.call(ctx, "tools/call", map[string]any{"name": name, "arguments": args}) + if err != nil { + return "", false, err + } + var reply struct { + Content []struct { + Type string `json:"type"` + Text *string `json:"text"` + } `json:"content"` + IsError *bool `json:"isError,omitempty"` + StructuredContent json.RawMessage `json:"structuredContent,omitempty"` + } + if acpDecode(result, &reply, false) != nil || reply.Content == nil { + return "", false, errACPMCP + } + for _, content := range reply.Content { + if content.Type != "text" || content.Text == nil { + return "", false, errACPMCP + } + } + if len(reply.StructuredContent) != 0 && reply.StructuredContent[0] != '{' { + return "", false, errACPMCP + } + // Only the validated model-visible projection crosses the provider + // boundary. MCP metadata and extension fields remain local to the client. + visible, err := json.Marshal(reply) + if err != nil { + return "", false, errACPMCP + } + return string(visible), reply.IsError != nil && *reply.IsError, nil +} + +func (m *acpMCPClient) call(ctx context.Context, method string, params any) (json.RawMessage, error) { + id, _ := json.Marshal(m.nextID.Add(1)) + encoded, err := json.Marshal(params) + if err != nil || len(encoded) > defaultMaxBrokeredBytes { + return nil, errACPMCP + } + body, _ := json.Marshal(acpRequest{JSONRPC: "2.0", ID: id, Method: method, Params: encoded}) + response, err := m.post(ctx, body) + if err != nil { + return nil, err + } + defer response.Body.Close() //nolint:errcheck + mediaType, _, err := mime.ParseMediaType(response.Header.Get("Content-Type")) + if err != nil || response.StatusCode != http.StatusOK || mediaType != "application/json" { + return nil, errACPMCP + } + data, err := io.ReadAll(io.LimitReader(response.Body, defaultMaxBrokeredBytes+1)) + if err != nil || len(data) > defaultMaxBrokeredBytes { + return nil, errACPMCP + } + var reply struct { + JSONRPC string `json:"jsonrpc"` + ID json.RawMessage `json:"id"` + Result json.RawMessage `json:"result"` + Error json.RawMessage `json:"error"` + } + if acpDecode(data, &reply, true) != nil || reply.JSONRPC != "2.0" || !bytes.Equal(reply.ID, id) || + len(reply.Error) != 0 || len(reply.Result) == 0 || reply.Result[0] != '{' { + return nil, errACPMCP + } + return reply.Result, nil +} + +func (m *acpMCPClient) post(ctx context.Context, body []byte) (*http.Response, error) { + request, err := http.NewRequestWithContext(ctx, http.MethodPost, m.url, bytes.NewReader(body)) + if err != nil { + return nil, errACPMCP + } + request.Header.Set("Authorization", m.bearer) + request.Header.Set("Content-Type", "application/json") + request.Header.Set("Accept", "application/json, text/event-stream") + request.Header.Set("MCP-Protocol-Version", acpMCPVersion) + response, err := m.client.Do(request) + if err != nil { + return nil, errACPMCP + } + return response, nil +} diff --git a/acp_mcp_metadata_test.go b/acp_mcp_metadata_test.go new file mode 100644 index 0000000..63683fa --- /dev/null +++ b/acp_mcp_metadata_test.go @@ -0,0 +1,77 @@ +package main + +import ( + "encoding/json" + "net/http" + "reflect" + "strings" + "sync/atomic" + "testing" +) + +func TestACPToolOutputForwardsOnlyValidatedModelContent(t *testing.T) { + for _, mode := range []string{"text", "structured", "error"} { + t.Run(mode, func(t *testing.T) { + content := []map[string]any{{"type": "text", "text": "héllo 世界"}} + want := map[string]any{"content": content} + if mode == "structured" { + // Structured content is model-visible data, including its own + // application fields. Only protocol metadata is omitted. + want["structuredContent"] = map[string]any{"value": "visible", "_meta": "application data"} + want["isError"] = false + } + if mode == "error" { + want["isError"] = true + } + mcp := &acpTestMCP{tools: func() []map[string]any { return acpTestTools("probe") }} + mcp.execute = func(w http.ResponseWriter, _ *http.Request, id json.RawMessage, _ string, _ json.RawMessage) { + result := map[string]any{} + for key, value := range want { + result[key] = value + } + result["content"] = []map[string]any{{ + "type": "text", "text": "héllo 世界", + "_meta": map[string]any{"private": "client-only-content-marker"}, + "extension": "unvalidated-content-marker", + }} + result["_meta"] = map[string]any{"private": "client-only-result-marker"} + result["extension"] = "unvalidated-result-marker" + acpTestMCPResult(w, id, result) + } + var requests atomic.Int32 + peer := newACPTestPeer(t, toolSchemaModeRequest, func(w http.ResponseWriter, r *http.Request) { + body := acpTestReadProvider(t, r) + switch requests.Add(1) { + case 1: + acpTestCompleted(w, "tool-response", "", acpTestCall("probe", "call-probe", `{}`)) + case 2: + var outputs []foundryFunctionOutput + if json.Unmarshal(body["input"], &outputs) != nil || len(outputs) != 1 { + t.Error("invalid function output continuation") + w.WriteHeader(http.StatusBadRequest) + return + } + var got, expected any + encoded, _ := json.Marshal(want) + if json.Unmarshal([]byte(outputs[0].Output), &got) != nil || json.Unmarshal(encoded, &expected) != nil || + !reflect.DeepEqual(got, expected) { + t.Error("provider function output was not limited to validated model content") + } + for _, marker := range []string{"client-only-", "unvalidated-"} { + if strings.Contains(outputs[0].Output, marker) { + t.Error("unvalidated tool metadata reached the provider") + } + } + acpTestCompleted(w, "final-response", "done") + default: + t.Error("unexpected provider request") + w.WriteHeader(http.StatusInternalServerError) + } + }, mcp) + acpAssertStop(t, peer.reply(peer.prompt("use probe")), "end_turn") + if requests.Load() != 2 || mcp.calls.Load() != 1 || acpOutput(peer.events) != "done" { + t.Fatal("tool continuation did not complete exactly once") + } + }) + } +} diff --git a/acp_prompt_test.go b/acp_prompt_test.go new file mode 100644 index 0000000..f19ffbc --- /dev/null +++ b/acp_prompt_test.go @@ -0,0 +1,455 @@ +package main + +import ( + "encoding/json" + "net/http" + "reflect" + "strings" + "sync/atomic" + "testing" + "time" +) + +func acpTestCall(name, id, arguments string) foundryOutputItem { + encoded, _ := json.Marshal(arguments) + return foundryOutputItem{Type: "function_call", CallID: id, Name: name, Arguments: encoded} +} + +func TestACPStdioSuccessfulContinuationAndLargeUnicodeOutput(t *testing.T) { + large := strings.Repeat("héllo 世界 🔒\n", 18000) + var requests atomic.Int32 + mcp := &acpTestMCP{} + peer := newACPTestPeer(t, toolSchemaModeRequest, func(w http.ResponseWriter, r *http.Request) { + body := acpTestReadProvider(t, r) + switch requests.Add(1) { + case 1: + if _, exists := body["previous_response_id"]; exists { + t.Error("fresh child replayed provider history") + } + acpTestCompleted(w, "opaque-first", large) + case 2: + if string(body["previous_response_id"]) != `"opaque-first"` { + t.Error("successful continuation lost previous response") + } + if string(body["input"]) != `"continue"` { + t.Error("continuation replayed earlier prompt content") + } + acpTestCompleted(w, "opaque-second", "done") + default: + t.Error("unexpected provider request") + w.WriteHeader(http.StatusInternalServerError) + } + }, mcp) + acpAssertStop(t, peer.reply(peer.prompt("first")), "end_turn") + if acpOutput(peer.events) != large { + t.Fatal("Unicode output was split incorrectly") + } + chunks := len(peer.events) + if chunks < 2 { + t.Fatal("large output was not bounded into multiple ACP frames") + } + for _, event := range peer.events { + if len(event["content"].(map[string]any)["text"].(string)) > acpTextChunkBytes { + t.Fatal("ACP chunk exceeded byte bound") + } + } + peer.events = nil + acpAssertStop(t, peer.reply(peer.prompt("continue")), "end_turn") + if acpOutput(peer.events) != "done" || mcp.lists.Load() != 2 || requests.Load() != 2 || mcp.calls.Load() != 0 { + t.Fatal("continuation did not use one fresh discovery and one provider request") + } +} + +func TestACPStaticToolsExactArgumentsConcurrentOutputAndFreshAllowlist(t *testing.T) { + const nested = `{"text":"héllo 世界 🌍","nested":{"items":["é","日本語",{"deeper":["🔒",null,false]}],"count":42,"enabled":true}}` + var requests atomic.Int32 + var revoked atomic.Bool + secondStarted := make(chan struct{}) + firstCompleted := make(chan struct{}) + mcp := &acpTestMCP{ + tools: func() []map[string]any { + if revoked.Load() { + return acpTestTools() + } + return acpTestTools("echo", "quick") + }, + execute: func(w http.ResponseWriter, r *http.Request, id json.RawMessage, name string, args json.RawMessage) { + switch name { + case "echo": + var want, got any + _ = json.Unmarshal([]byte(nested), &want) + _ = json.Unmarshal(args, &got) + if !reflect.DeepEqual(want, got) { + t.Error("nested Unicode tool arguments changed") + } + select { + case <-secondStarted: + case <-r.Context().Done(): + return + } + acpTestToolResult(w, id, string(args), false) + close(firstCompleted) + case "quick": + close(secondStarted) + acpTestToolResult(w, id, `{"quick":true}`, false) + default: + t.Error("unlisted tool executed") + } + }, + } + peer := newACPTestPeer(t, toolSchemaModeProviderStatic, func(w http.ResponseWriter, r *http.Request) { + body := acpTestReadProvider(t, r) + if _, exists := body["tools"]; exists { + t.Error("static mode sent request-level tool schemas") + } + switch requests.Add(1) { + case 1: + acpTestCompleted(w, "tools-response", "", acpTestCall("echo", "provider-call-private-1", nested), acpTestCall("quick", "provider-call-private-2", `{}`)) + case 2: + select { + case <-firstCompleted: + default: + t.Error("provider resumed before all tool calls joined") + } + var outputs []foundryFunctionOutput + if json.Unmarshal(body["input"], &outputs) != nil || len(outputs) != 2 { + t.Error("invalid function output continuation") + } + for i, output := range outputs { + if output.Type != "function_call_output" || output.CallID != []string{"provider-call-private-1", "provider-call-private-2"}[i] { + t.Error("function output correlation changed") + } + } + if string(body["previous_response_id"]) != `"tools-response"` { + t.Error("tool continuation lost provider checkpoint") + } + acpTestCompleted(w, "tools-final", "finished") + case 3: + if string(body["previous_response_id"]) != `"tools-final"` { + t.Error("prompt continuation lost final checkpoint") + } + acpTestCompleted(w, "forbidden-call", "", acpTestCall("echo", "forbidden", `{}`)) + default: + t.Error("provider call replayed") + w.WriteHeader(http.StatusInternalServerError) + } + }, mcp) + acpAssertStop(t, peer.reply(peer.prompt("nested batch")), "end_turn") + acpAssertToolEvents(t, peer.events, 2, 0) + encoded, _ := json.Marshal(peer.events) + if strings.Contains(string(encoded), "provider-call-private") || strings.Contains(string(encoded), "héllo") || strings.Contains(string(encoded), "test-only") { + t.Fatal("tool lifecycle exposed provider IDs, arguments, or credentials") + } + revoked.Store(true) + peer.events = nil + acpAssertFailure(t, peer.reply(peer.prompt("fresh allowlist"))) + if mcp.lists.Load() != 2 || mcp.calls.Load() != 2 || requests.Load() != 3 || len(peer.events) != 0 { + t.Fatal("stale allowed tool was executed or replayed") + } +} + +func TestACPMCPAdmittedErrorRemainsRecoverable(t *testing.T) { + var requests atomic.Int32 + mcp := &acpTestMCP{tools: func() []map[string]any { return acpTestTools("probe") }} + mcp.execute = func(w http.ResponseWriter, _ *http.Request, id json.RawMessage, _ string, _ json.RawMessage) { + failed := mcp.calls.Load() == 1 + acpTestToolResult(w, id, `{"status":"fixture"}`, failed) + } + peer := newACPTestPeer(t, toolSchemaModeRequest, func(w http.ResponseWriter, r *http.Request) { + body := acpTestReadProvider(t, r) + if _, present := body["tools"]; !present { + t.Error("request schema mode omitted discovered tool") + } + switch requests.Add(1) { + case 1: + acpTestCompleted(w, "response-a", "", acpTestCall("probe", "call-a", `{}`)) + case 2: + var outputs []foundryFunctionOutput + if json.Unmarshal(body["input"], &outputs) != nil || len(outputs) != 1 || !strings.Contains(outputs[0].Output, `"isError":true`) { + t.Error("admitted error did not reach model as function output") + } + acpTestCompleted(w, "response-b", "", acpTestCall("probe", "call-b", `{}`)) + case 3: + acpTestCompleted(w, "response-c", "recovered") + default: + t.Error("unexpected model retry") + w.WriteHeader(http.StatusInternalServerError) + } + }, mcp) + acpAssertStop(t, peer.reply(peer.prompt("recover")), "end_turn") + acpAssertToolEvents(t, peer.events, 1, 1) + if requests.Load() != 3 || mcp.calls.Load() != 2 || acpOutput(peer.events) != "recovered" { + t.Fatal("admitted tool recovery failed") + } +} + +func TestACPFatalMCPFailureCancelsAndJoinsSiblingWithoutContinuation(t *testing.T) { + for _, failure := range []string{"rpc", "http", "malformed", "wrong-id", "wrong-version", "both-result-error"} { + t.Run(failure, func(t *testing.T) { + var requests atomic.Int32 + slowStarted := make(chan struct{}) + slowCancelled := make(chan struct{}) + mcp := &acpTestMCP{tools: func() []map[string]any { return acpTestTools("probe") }} + mcp.execute = func(w http.ResponseWriter, r *http.Request, id json.RawMessage, _ string, args json.RawMessage) { + var params struct { + Slot int `json:"slot"` + } + _ = json.Unmarshal(args, ¶ms) + if params.Slot == 1 { + close(slowStarted) + <-r.Context().Done() + close(slowCancelled) + return + } + select { + case <-slowStarted: + case <-r.Context().Done(): + return + } + w.Header().Set("Content-Type", "application/json") + switch failure { + case "rpc": + _ = json.NewEncoder(w).Encode(map[string]any{"jsonrpc": "2.0", "id": id, "error": map[string]any{"code": -32002, "message": "test-only-sensitive-provider-detail"}}) + case "http": + w.WriteHeader(http.StatusServiceUnavailable) + case "malformed": + _, _ = w.Write([]byte(`{"jsonrpc":`)) + case "wrong-id": + acpTestToolResult(w, json.RawMessage("9999"), "not correlated", false) + case "wrong-version": + _ = json.NewEncoder(w).Encode(map[string]any{"jsonrpc": "1.0", "id": id, "result": map[string]any{}}) + case "both-result-error": + _ = json.NewEncoder(w).Encode(map[string]any{"jsonrpc": "2.0", "id": id, "result": map[string]any{}, "error": map[string]any{"code": -1}}) + } + } + peer := newACPTestPeer(t, toolSchemaModeProviderStatic, func(w http.ResponseWriter, r *http.Request) { + acpTestReadProvider(t, r) + requests.Add(1) + acpTestCompleted(w, "fatal-batch", "", acpTestCall("probe", "private-fatal", `{"slot":0}`), acpTestCall("probe", "private-slow", `{"slot":1}`)) + }, mcp) + acpAssertFailure(t, peer.reply(peer.prompt("fatal batch"))) + select { + case <-slowCancelled: + case <-time.After(2 * time.Second): + t.Fatal("fatal MCP failure left sibling request active") + } + acpAssertToolEvents(t, peer.events, 0, 2) + if requests.Load() != 1 || mcp.calls.Load() != 2 || acpOutput(peer.events) != "" { + t.Fatal("fatal MCP failure continued or replayed") + } + encoded, _ := json.Marshal(peer.events) + if strings.Contains(string(encoded), "private-") || strings.Contains(string(encoded), "test-only-") { + t.Fatal("fatal tool metadata leaked sensitive detail") + } + if peer.reply(peer.prompt("must not reuse poisoned child"))["error"] == nil || requests.Load() != 1 { + t.Fatal("failed child resumed provider state") + } + }) + } +} + +func TestACPCancelJoinsToolsUnderStdoutBackpressure(t *testing.T) { + for _, method := range []string{"session/cancel", "$/cancel_request"} { + t.Run(method, func(t *testing.T) { + var requests atomic.Int32 + started := make(chan struct{}, 2) + cancelled := make(chan struct{}, 2) + mcp := &acpTestMCP{ + tools: func() []map[string]any { return acpTestTools("hold") }, + execute: func(_ http.ResponseWriter, r *http.Request, _ json.RawMessage, _ string, _ json.RawMessage) { + started <- struct{}{} + <-r.Context().Done() + cancelled <- struct{}{} + }, + } + peer := newACPTestPeer(t, toolSchemaModeRequest, func(w http.ResponseWriter, r *http.Request) { + acpTestReadProvider(t, r) + requests.Add(1) + acpTestCompleted(w, "held", "", acpTestCall("hold", "hold-a", `{}`), acpTestCall("hold", "hold-b", `{}`)) + }, mcp) + id := peer.prompt("hold two calls") + for range 2 { + message := peer.read() + if message["method"] != "session/update" { + t.Fatal("missing held tool start") + } + } + for range 2 { + select { + case <-started: + case <-time.After(2 * time.Second): + t.Fatal("tools did not overlap") + } + } + params := map[string]any{"sessionId": peer.session} + if method == "$/cancel_request" { + params = map[string]any{"requestId": id} + } + peer.send(map[string]any{"jsonrpc": "2.0", "method": method, "params": params}) + // Deliberately do not read stdout until both HTTP contexts are gone. + for range 2 { + select { + case <-cancelled: + case <-time.After(2 * time.Second): + t.Fatal("stdout pressure blocked tool cancellation") + } + } + acpAssertStop(t, peer.reply(id), "cancelled") + acpAssertToolEvents(t, peer.events, 0, 2) + if requests.Load() != 1 || mcp.calls.Load() != 2 { + t.Fatal("cancelled batch was replayed") + } + }) + } +} + +func TestACPFatalMCPFailureCancelsSiblingBeforeBlockedEventWrite(t *testing.T) { + var requests atomic.Int32 + started, cancelled := make(chan struct{}), make(chan struct{}) + releaseFailure := make(chan struct{}) + mcp := &acpTestMCP{tools: func() []map[string]any { return acpTestTools("probe") }} + mcp.execute = func(w http.ResponseWriter, r *http.Request, id json.RawMessage, _ string, args json.RawMessage) { + var params struct { + Slot int `json:"slot"` + } + _ = json.Unmarshal(args, ¶ms) + if params.Slot == 1 { + close(started) + <-r.Context().Done() + close(cancelled) + return + } + select { + case <-releaseFailure: + case <-r.Context().Done(): + return + } + w.Header().Set("Content-Type", "application/json") + _ = json.NewEncoder(w).Encode(map[string]any{"jsonrpc": "2.0", "id": id, + "error": map[string]any{"code": -32002, "message": "test-only-protocol-failure"}}) + } + peer := newACPTestPeer(t, toolSchemaModeRequest, func(w http.ResponseWriter, r *http.Request) { + acpTestReadProvider(t, r) + requests.Add(1) + acpTestCompleted(w, "fatal-paused", "", acpTestCall("probe", "fatal", `{"slot":0}`), acpTestCall("probe", "held", `{"slot":1}`)) + }, mcp) + id := peer.prompt("fatal failure with stdout paused") + for range 2 { + if peer.read()["method"] != "session/update" { + t.Fatal("missing tool start") + } + } + select { + case <-started: + case <-time.After(2 * time.Second): + t.Fatal("held sibling did not start") + } + close(releaseFailure) + // A terminal event cannot be delivered until read resumes. Its delivery + // must not be a prerequisite for revoking the still-running HTTP call. + select { + case <-cancelled: + case <-time.After(2 * time.Second): + t.Fatal("fatal MCP error left sibling running behind blocked stdout") + } + acpAssertFailure(t, peer.reply(id)) + acpAssertToolEvents(t, peer.events, 0, 2) + if requests.Load() != 1 || mcp.calls.Load() != 2 || acpOutput(peer.events) != "" { + t.Fatal("fatal failure continued the model or replayed tool calls") + } +} + +func TestACPProviderFailureAfterToolDoesNotReplayOrCommitOutput(t *testing.T) { + var requests atomic.Int32 + mcp := &acpTestMCP{ + tools: func() []map[string]any { return acpTestTools("probe") }, + execute: func(w http.ResponseWriter, _ *http.Request, id json.RawMessage, _ string, _ json.RawMessage) { + acpTestToolResult(w, id, "side effect completed", false) + }, + } + peer := newACPTestPeer(t, toolSchemaModeRequest, func(w http.ResponseWriter, r *http.Request) { + acpTestReadProvider(t, r) + if requests.Add(1) == 1 { + acpTestCompleted(w, "before-failure", "uncommitted draft", acpTestCall("probe", "once", `{}`)) + } else { + w.WriteHeader(http.StatusServiceUnavailable) + _, _ = w.Write([]byte("test-only-sensitive-provider-detail")) + } + }, mcp) + acpAssertFailure(t, peer.reply(peer.prompt("provider failure"))) + acpAssertToolEvents(t, peer.events, 1, 0) + if requests.Load() != 2 || mcp.calls.Load() != 1 || acpOutput(peer.events) != "" { + t.Fatal("provider fault replayed tool or committed draft output") + } +} + +func TestACPRejectsUnsupportedPromptAndSessionCapabilities(t *testing.T) { + var requests atomic.Int32 + peer := newACPTestPeer(t, toolSchemaModeRequest, func(w http.ResponseWriter, r *http.Request) { + acpTestReadProvider(t, r) + requests.Add(1) + acpTestCompleted(w, "valid", "ok") + }, &acpTestMCP{}) + for _, method := range []string{"session/load", "session/request_permission", "fs/read_text_file", "terminal/create", "authenticate"} { + if peer.call(method, map[string]any{})["error"] == nil { + t.Fatal("unsupported capability accepted") + } + } + if peer.call("session/new", map[string]any{})["error"] == nil { + t.Fatal("second child session accepted") + } + if peer.call("session/prompt", map[string]any{"sessionId": peer.session, "prompt": []map[string]string{{"type": "image", "data": "unsupported"}}})["error"] == nil { + t.Fatal("unsupported prompt block accepted") + } + if requests.Load() != 0 { + t.Fatal("unsupported request reached provider") + } + acpAssertStop(t, peer.reply(peer.prompt("valid")), "end_turn") +} + +func TestACPCancellationClosesProviderStream(t *testing.T) { + started, cancelled := make(chan struct{}), make(chan struct{}) + var requests atomic.Int32 + peer := newACPTestPeer(t, toolSchemaModeRequest, func(w http.ResponseWriter, r *http.Request) { + acpTestReadProvider(t, r) + requests.Add(1) + w.Header().Set("Content-Type", "text/event-stream") + _, _ = w.Write([]byte(acpTestSSE(`{"type":"response.created","response":{"id":"held","status":"in_progress"}}`, `{"type":"response.output_text.delta","delta":"uncommitted"}`))) + w.(http.Flusher).Flush() + close(started) + <-r.Context().Done() + close(cancelled) + }, &acpTestMCP{}) + id := peer.prompt("hold provider") + select { + case <-started: + case <-time.After(2 * time.Second): + t.Fatal("provider did not start") + } + peer.send(map[string]any{"jsonrpc": "2.0", "method": "session/cancel", "params": map[string]string{"sessionId": peer.session}}) + acpAssertStop(t, peer.reply(id), "cancelled") + select { + case <-cancelled: + case <-time.After(2 * time.Second): + t.Fatal("provider connection remained active after cancellation") + } + if requests.Load() != 1 || len(peer.events) != 0 { + t.Fatal("cancelled stream replayed or committed text") + } +} + +func TestACPResourceLinkProjectsTextWithoutFilesystemOrHTTPAccess(t *testing.T) { + peer := newACPTestPeer(t, toolSchemaModeRequest, func(w http.ResponseWriter, r *http.Request) { + body := acpTestReadProvider(t, r) + var input string + if json.Unmarshal(body["input"], &input) != nil || input != "inspect\nResource link: file\nURI: file:///workspace/file\nMIME type: text/plain" { + t.Error("resource link projection changed") + } + acpTestCompleted(w, "resource", "ok") + }, &acpTestMCP{}) + response := peer.call("session/prompt", map[string]any{"sessionId": peer.session, "prompt": []map[string]string{ + {"type": "text", "text": "inspect"}, + {"type": "resource_link", "name": "file", "uri": "file:///workspace/file", "mimeType": "text/plain"}, + }}) + acpAssertStop(t, response, "end_turn") +} diff --git a/acp_protocol.go b/acp_protocol.go new file mode 100644 index 0000000..f62d964 --- /dev/null +++ b/acp_protocol.go @@ -0,0 +1,274 @@ +package main + +import ( + "bufio" + "bytes" + "encoding/json" + "errors" + "io" + "reflect" + "strconv" + "strings" + "unicode/utf8" +) + +type acpRPCError struct { + Code int `json:"code"` + Message string `json:"message"` +} + +func (e *acpRPCError) Error() string { return e.Message } + +var ( + acpInvalidRequest = &acpRPCError{-32600, "invalid ACP request"} + acpInvalidParams = &acpRPCError{-32602, "invalid ACP parameters"} + acpInternalError = &acpRPCError{-32603, "Foundry ACP prompt failed"} +) + +type acpRequest struct { + JSONRPC string `json:"jsonrpc"` + ID json.RawMessage `json:"id,omitempty"` + Method string `json:"method"` + Params json.RawMessage `json:"params,omitempty"` +} + +type acpResponse struct { + JSONRPC string `json:"jsonrpc"` + ID json.RawMessage `json:"id"` + Result any `json:"result,omitempty"` + Error *acpRPCError `json:"error,omitempty"` +} + +type acpMCPServer struct { + Type string `json:"type"` + Name string `json:"name"` + URL string `json:"url"` + Headers []struct { + Name string `json:"name"` + Value string `json:"value"` + } `json:"headers"` + Meta json.RawMessage `json:"_meta,omitempty"` +} + +type acpNewSession struct { + CWD string `json:"cwd"` + AdditionalDirectories []string `json:"additionalDirectories,omitempty"` + MCPServers []acpMCPServer `json:"mcpServers"` + Meta json.RawMessage `json:"_meta,omitempty"` +} + +type acpPrompt struct { + SessionID string `json:"sessionId"` + Prompt []struct { + Type string `json:"type"` + Text *string `json:"text,omitempty"` + Name string `json:"name,omitempty"` + URI string `json:"uri,omitempty"` + MIMEType string `json:"mimeType,omitempty"` + Meta json.RawMessage `json:"_meta,omitempty"` + } `json:"prompt"` + Meta json.RawMessage `json:"_meta,omitempty"` +} + +func (p acpPrompt) text() (string, error) { + if !acpSafeString(p.SessionID, 512) || len(p.Prompt) == 0 { + return "", acpInvalidParams + } + var blocks []string + for _, block := range p.Prompt { + switch block.Type { + case "text": + if block.Text == nil || block.Name != "" || block.URI != "" || block.MIMEType != "" { + return "", acpInvalidParams + } + blocks = append(blocks, *block.Text) + case "resource_link": + if block.Text != nil || !acpSafeString(block.Name, 1024) || !acpSafeString(block.URI, 8<<10) { + return "", acpInvalidParams + } + text := "Resource link: " + block.Name + "\nURI: " + block.URI + if block.MIMEType != "" { + if !acpSafeString(block.MIMEType, 256) { + return "", acpInvalidParams + } + text += "\nMIME type: " + block.MIMEType + } + blocks = append(blocks, text) + default: + return "", acpInvalidParams + } + } + text := strings.Join(blocks, "\n") + if len(text) > maxFoundryPromptBytes { + return "", acpInvalidParams + } + return text, nil +} + +func acpRequestKey(raw json.RawMessage) (string, error) { + if len(raw) == 0 || len(raw) > 1024 { + return "", acpInvalidRequest + } + if raw[0] == '"' { + var value string + if json.Unmarshal(raw, &value) != nil || !acpSafeString(value, 512) { + return "", acpInvalidRequest + } + return "s:" + value, nil + } + value, err := strconv.ParseInt(string(raw), 10, 64) + if err != nil { + return "", acpInvalidRequest + } + return "n:" + strconv.FormatInt(value, 10), nil +} + +func acpReadLine(reader *bufio.Reader) ([]byte, error) { + var line []byte + for { + part, err := reader.ReadSlice('\n') + if len(line)+len(part) > acpMaxMessageBytes { + return nil, acpInvalidRequest + } + line = append(line, part...) + if err == nil { + return bytes.TrimSpace(line), nil + } + if errors.Is(err, bufio.ErrBufferFull) { + continue + } + if len(line) != 0 { + return nil, acpInvalidRequest + } + return nil, err + } +} + +// Go's ordinary JSON decoder accepts duplicate object members. Reject them at +// every depth before interpreting authority-bearing envelopes or tool arguments. +func acpDecode(data []byte, value any, strictFields bool) error { + return acpDecodeJSON(data, value, strictFields, nil) +} + +// Remote authority and Responses structs accept single case-folded field names, +// as encoding/json does, but two names must never replace or merge one field. +// Maps, interfaces and custom JSON values remain opaque to field folding. +func acpDecodeStruct(data []byte, value any, strictFields bool) error { + return acpDecodeJSON(data, value, strictFields, reflect.TypeOf(value)) +} + +func acpDecodeJSON(data []byte, value any, strictFields bool, shape reflect.Type) error { + if !utf8.Valid(data) || !json.Valid(data) || !acpValidStringEscapes(data) { + return acpInvalidParams + } + check := json.NewDecoder(bytes.NewReader(data)) + check.UseNumber() + if acpJSONValue(check, 0, shape) != nil { + return acpInvalidParams + } + if _, err := check.Token(); !errors.Is(err, io.EOF) { + return acpInvalidParams + } + decoder := json.NewDecoder(bytes.NewReader(data)) + decoder.UseNumber() + if strictFields { + decoder.DisallowUnknownFields() + } + if err := decoder.Decode(value); err != nil { + return acpInvalidParams + } + return nil +} + +// encoding/json replaces unpaired UTF-16 escapes with U+FFFD. Tool arguments +// must retain their exact Unicode value, so reject malformed surrogate pairs. +func acpValidStringEscapes(data []byte) bool { + quoted := false + for i := 0; i < len(data); i++ { + if data[i] == '"' { + quoted = !quoted + continue + } + if !quoted || data[i] != '\\' { + continue + } + i++ + if data[i] != 'u' { + continue + } + value, err := strconv.ParseUint(string(data[i+1:i+5]), 16, 16) + if err != nil { + return false + } + i += 4 + if value >= 0xdc00 && value <= 0xdfff { + return false + } + if value < 0xd800 || value > 0xdbff { + continue + } + if i+6 >= len(data) || data[i+1] != '\\' || data[i+2] != 'u' { + return false + } + low, err := strconv.ParseUint(string(data[i+3:i+7]), 16, 16) + if err != nil || low < 0xdc00 || low > 0xdfff { + return false + } + i += 6 + } + return true +} + +func acpJSONValue(decoder *json.Decoder, depth int, shape reflect.Type) error { + if depth > 64 { + return acpInvalidParams + } + token, err := decoder.Token() + if err != nil { + return err + } + delim, ok := token.(json.Delim) + if !ok { + return nil + } + switch delim { + case '{': + fields, err := acpJSONStructFields(shape, 0) + if err != nil { + return err + } + seen := make(map[string]bool) + for decoder.More() { + key, err := decoder.Token() + if err != nil { + return err + } + name, ok := key.(string) + if !ok { + return acpInvalidParams + } + name, fieldType := acpJSONMatchField(fields, name) + if seen[name] { + return acpInvalidParams + } + seen[name] = true + if err := acpJSONValue(decoder, depth+1, fieldType); err != nil { + return err + } + } + case '[': + var element reflect.Type + if shape := acpJSONStructType(shape); shape != nil && (shape.Kind() == reflect.Slice || shape.Kind() == reflect.Array) { + element = shape.Elem() + } + for decoder.More() { + if err := acpJSONValue(decoder, depth+1, element); err != nil { + return err + } + } + default: + return acpInvalidParams + } + _, err = decoder.Token() + return err +} diff --git a/acp_protocol_test.go b/acp_protocol_test.go new file mode 100644 index 0000000..becbb4d --- /dev/null +++ b/acp_protocol_test.go @@ -0,0 +1,88 @@ +package main + +import ( + "bufio" + "errors" + "io" + "net/http" + "strings" + "sync/atomic" + "testing" +) + +func TestACPReadLineRequiresBoundedNewlineFrames(t *testing.T) { + for _, test := range []struct { + name string + input string + want string + err error + }{ + {name: "LF", input: " {\"id\":1}\n", want: `{"id":1}`}, + {name: "CRLF", input: "\t{\"id\":1}\r\n", want: `{"id":1}`}, + {name: "blank", input: " \t\r\n"}, + {name: "empty EOF", err: io.EOF}, + {name: "unterminated object", input: `{"id":1}`, err: acpInvalidRequest}, + {name: "unterminated whitespace", input: " ", err: acpInvalidRequest}, + {name: "exact bound", input: strings.Repeat("x", acpMaxMessageBytes-1) + "\n", want: strings.Repeat("x", acpMaxMessageBytes-1)}, + {name: "over bound", input: strings.Repeat("x", acpMaxMessageBytes) + "\n", err: acpInvalidRequest}, + } { + t.Run(test.name, func(t *testing.T) { + // A small reader forces the multi-fragment path independently of + // the operating system's pipe scheduling. + reader := bufio.NewReaderSize(strings.NewReader(test.input), 31) + line, err := acpReadLine(reader) + if !errors.Is(err, test.err) || string(line) != test.want { + t.Fatal("newline framing or its byte limit changed") + } + }) + } + reader := bufio.NewReaderSize(strings.NewReader("first\nsecond\n"), 31) + for _, want := range []string{"first", "second"} { + if line, err := acpReadLine(reader); err != nil || string(line) != want { + t.Fatal("adjacent newline frames were combined or lost") + } + } + if _, err := acpReadLine(reader); !errors.Is(err, io.EOF) { + t.Fatal("complete frames did not terminate with EOF") + } +} + +func TestACPInvalidEnvelopesCannotReachProvider(t *testing.T) { + var requests atomic.Int32 + peer := newACPTestPeer(t, toolSchemaModeRequest, func(w http.ResponseWriter, r *http.Request) { + acpTestReadProvider(t, r) + requests.Add(1) + acpTestCompleted(w, "valid", "ok") + }, &acpTestMCP{}) + for _, test := range []struct { + input string + code float64 + }{ + {`{`, -32700}, + {`[]`, -32600}, + {`null`, -32600}, + {`{"jsonrpc":"1.0","id":100,"method":"initialize"}`, -32600}, + {`{"jsonrpc":"2.0","id":null,"method":"initialize"}`, -32600}, + {`{"jsonrpc":"2.0","id":true,"method":"initialize"}`, -32600}, + {`{"jsonrpc":"2.0","id":1.5,"method":"initialize"}`, -32600}, + {`{"jsonrpc":"2.0","id":100,"method":"initialize","result":{}}`, -32600}, + {`{"jsonrpc":"2.0","id":100,"method":"initialize","params":{"a":1,"a":2}}`, -32600}, + {`{"jsonrpc":"2.0","id":100,"id":101,"method":"initialize"}`, -32600}, + } { + if _, err := io.WriteString(peer.in, test.input+"\n"); err != nil { + t.Fatal("could not send invalid fixture envelope") + } + reply := peer.read() + failure, ok := reply["error"].(map[string]any) + if !ok || failure["code"] != test.code || reply["id"] != nil || reply["result"] != nil { + t.Fatal("malformed envelope did not produce an uncorrelated protocol error") + } + } + if requests.Load() != 0 { + t.Fatal("invalid envelope invoked provider") + } + acpAssertStop(t, peer.reply(peer.prompt("valid after malformed envelopes")), "end_turn") + if requests.Load() != 1 { + t.Fatal("malformed envelope changed a valid session") + } +} diff --git a/acp_responses.go b/acp_responses.go new file mode 100644 index 0000000..b600d08 --- /dev/null +++ b/acp_responses.go @@ -0,0 +1,280 @@ +package main + +import ( + "bufio" + "bytes" + "context" + "encoding/json" + "io" + "mime" + "net/http" + "strings" +) + +type acpResponseRequest struct { + foundryResponseRequest + Model string `json:"model"` +} + +// ACP never constructs a Foundry SDK client: the privileged supervisor owns the +// remote agent session and rewrites that binding outside this child process. +func acpCreateResponse(ctx context.Context, cfg acpConfiguration, client *http.Client, request foundryResponseRequest) (foundryStreamSummary, error) { + request.Stream, request.Store = true, true + request.AgentSessionID = "" + body, err := json.Marshal(acpResponseRequest{foundryResponseRequest: request, Model: cfg.agent.Model}) + if err != nil { + return foundryStreamSummary{}, errACPProvider + } + httpRequest, err := http.NewRequestWithContext(ctx, http.MethodPost, cfg.providerURL, bytes.NewReader(body)) + if err != nil { + return foundryStreamSummary{}, errACPProvider + } + httpRequest.Header.Set("Authorization", "Bearer "+cfg.token) + httpRequest.Header.Set("Content-Type", "application/json") + httpRequest.Header.Set("Accept", "text/event-stream, application/json") + response, err := client.Do(httpRequest) + if err != nil { + return foundryStreamSummary{}, errACPProvider + } + defer response.Body.Close() //nolint:errcheck + mediaType, _, err := mime.ParseMediaType(response.Header.Get("Content-Type")) + if err != nil || response.StatusCode != http.StatusOK { + return foundryStreamSummary{}, errACPProvider + } + switch mediaType { + case "text/event-stream": + return acpParseFoundrySSE(response.Body) + case "application/json": + data, err := io.ReadAll(io.LimitReader(response.Body, defaultMaxStreamBytes+1)) + if err != nil || len(data) > defaultMaxStreamBytes { + return foundryStreamSummary{}, errACPProvider + } + document, err := acpDecodeFoundryResponse(data) + if err != nil || document.Status != "completed" { + return foundryStreamSummary{}, errACPProvider + } + summary, err := processCompletedResponse(document, responseCallbacks{}) + if err != nil || acpValidateSummary(summary) != nil { + return foundryStreamSummary{}, errACPProvider + } + return summary, nil + default: + return foundryStreamSummary{}, errACPProvider + } +} + +func acpDecodeFoundryResponse(data []byte) (foundryResponse, error) { + var response foundryResponse + if acpDecodeStruct(data, &response, false) != nil || response.ID == "" || validateProviderIdentifier("response", response.ID) != nil || + response.Error != nil || response.Incomplete != nil { + return foundryResponse{}, errACPProvider + } + var fields map[string]json.RawMessage + if json.Unmarshal(data, &fields) != nil { + return foundryResponse{}, errACPProvider + } + var rawOutput json.RawMessage + for name, value := range fields { + if strings.EqualFold(name, "output") { + if rawOutput != nil { + return foundryResponse{}, errACPProvider + } + rawOutput = value + } + } + var output []json.RawMessage + if rawOutput != nil && json.Unmarshal(rawOutput, &output) != nil { + return foundryResponse{}, errACPProvider + } + // Return the validated objects, never a separately decoded typed slice + // whose elements could retain fields across folded output aliases. + response.Output = nil + for _, rawItem := range output { + item, err := acpDecodeFoundryItem(rawItem, true) + if err != nil { + return foundryResponse{}, err + } + response.Output = append(response.Output, item) + } + return response, nil +} + +func acpDecodeFoundryItem(data []byte, done bool) (foundryOutputItem, error) { + var item struct { + foundryOutputItem + Status string `json:"status"` + Role string `json:"role"` + } + if acpDecodeStruct(data, &item, false) != nil || (done && item.Status != "" && item.Status != "completed") { + return foundryOutputItem{}, errACPProvider + } + switch item.Type { + case "message": + if item.Role != "" && item.Role != "assistant" { + return foundryOutputItem{}, errACPProvider + } + for _, content := range item.Content { + if content.Type != "output_text" { + return foundryOutputItem{}, errACPProvider + } + } + case "function_call", "reasoning": + default: + // Hosted/native tools have no authority in the child. Only ordinary + // function calls, later checked against tools/list, may execute. + return foundryOutputItem{}, errACPProvider + } + return item.foundryOutputItem, nil +} + +func acpValidateSummary(summary foundryStreamSummary) error { + if summary.Status != "completed" || summary.ResponseID == "" || + validateProviderIdentifier("response", summary.ResponseID) != nil || + summary.Error != nil || summary.Incomplete != nil || len(summary.Text) > defaultMaxOutputBytes || + len(summary.FunctionCalls) > defaultMaxBrokeredCalls { + return errACPProvider + } + return nil +} + +// Keep the existing Responses types and terminal-output reconciliation, but +// require one explicit, coherent terminal event. A created/in-progress status, +// [DONE] alone, or a complete function-call item does not settle a response. +func acpParseFoundrySSE(reader io.Reader) (foundryStreamSummary, error) { + limited := &io.LimitedReader{R: reader, N: defaultMaxStreamBytes + 1} + scanner := bufio.NewScanner(limited) + scanner.Buffer(make([]byte, 32<<10), defaultMaxEventBytes) + var summary foundryStreamSummary + var data []byte + terminal, done := false, false + events := 0 + pending := map[string]bool{} + apply := func() error { + if len(data) == 0 { + return nil + } + events++ + if events > defaultMaxEvents || done { + return errACPProvider + } + if bytes.Equal(bytes.TrimSpace(data), []byte("[DONE]")) { + if !terminal { + return errACPProvider + } + done = true + return nil + } + if terminal { + return errACPProvider + } + var event foundryResponseEvent + var rawFields map[string]json.RawMessage + if acpDecodeStruct(data, &event, false) != nil || json.Unmarshal(data, &rawFields) != nil { + return errACPProvider + } + // Match the struct decoder's Unicode field folding without allowing two + // envelope members to validate one value and apply another. Tool argument + // objects keep their case-sensitive keys. + fields := make(map[string]json.RawMessage, 7) + for key, value := range rawFields { + for _, name := range []string{"type", "delta", "sequence_number", "response", "item", "error", "item_id"} { + if strings.EqualFold(key, name) { + if _, exists := fields[name]; exists { + return errACPProvider + } + fields[name] = value + break + } + } + } + if event.Response != nil { + response, err := acpDecodeFoundryResponse(fields["response"]) + if err != nil || (summary.ResponseID != "" && summary.ResponseID != response.ID) { + return errACPProvider + } + event.Response = &response + } + switch event.Type { + case "response.created", "response.in_progress", "response.queued": + if event.Response == nil || (event.Response.Status != "in_progress" && event.Response.Status != "queued") { + return errACPProvider + } + case "response.completed": + if event.Response == nil || event.Response.Status != "completed" { + return errACPProvider + } + for _, item := range event.Response.Output { + if item.Type == "function_call" { + delete(pending, item.ID) + } + } + if len(pending) != 0 { + return errACPProvider + } + terminal = true + case "response.output_item.added", "response.output_item.done": + item, err := acpDecodeFoundryItem(fields["item"], event.Type == "response.output_item.done") + if err != nil { + return err + } + event.Item = &item + if item.Type == "function_call" { + if event.Type == "response.output_item.added" { + if !acpSafeString(item.ID, maxProviderIdentifierBytes) { + return errACPProvider + } + pending[item.ID] = true + } else { + delete(pending, item.ID) + } + } + case "response.function_call_arguments.delta", "response.function_call_arguments.done": + var itemID string + if json.Unmarshal(fields["item_id"], &itemID) != nil || !acpSafeString(itemID, maxProviderIdentifierBytes) { + return errACPProvider + } + pending[itemID] = true + case "response.output_text.delta": + if len(fields["delta"]) == 0 || fields["delta"][0] != '"' || len(summary.Text)+len(event.Delta) > defaultMaxOutputBytes { + return errACPProvider + } + case "response.output_text.done", "response.content_part.added", "response.content_part.done", + "response.reasoning_summary_text.delta", "response.reasoning_summary_text.done", + "response.reasoning_summary_part.added", "response.reasoning_summary_part.done", + "response.reasoning_text.delta", "response.reasoning_text.done": + default: + return errACPProvider + } + if err := applyFoundryEvent(&summary, event, responseCallbacks{}); err != nil || len(summary.Text) > defaultMaxOutputBytes { + return errACPProvider + } + return nil + } + for scanner.Scan() { + line := scanner.Bytes() + if len(line) == 0 { + if err := apply(); err != nil { + return foundryStreamSummary{}, err + } + data = nil + continue + } + if part, ok := bytes.CutPrefix(line, []byte("data:")); ok { + part = bytes.TrimPrefix(part, []byte(" ")) + if len(data)+len(part)+1 > defaultMaxEventBytes { + return foundryStreamSummary{}, errACPProvider + } + if len(data) != 0 { + data = append(data, '\n') + } + data = append(data, part...) + } else if !bytes.HasPrefix(line, []byte(":")) && !strings.HasPrefix(string(line), "event:") && + !strings.HasPrefix(string(line), "id:") && !strings.HasPrefix(string(line), "retry:") { + return foundryStreamSummary{}, errACPProvider + } + } + if scanner.Err() != nil || limited.N <= 0 || len(data) != 0 || !terminal || acpValidateSummary(summary) != nil { + return foundryStreamSummary{}, errACPProvider + } + return summary, nil +} diff --git a/acp_responses_test.go b/acp_responses_test.go new file mode 100644 index 0000000..1a2c899 --- /dev/null +++ b/acp_responses_test.go @@ -0,0 +1,310 @@ +package main + +import ( + "encoding/json" + "fmt" + "net/http" + "strings" + "sync/atomic" + "testing" +) + +func acpTestSSE(events ...string) string { + return "data: " + strings.Join(events, "\n\ndata: ") + "\n\n" +} + +func TestACPResponsesRequireExplicitCoherentCompletion(t *testing.T) { + const created = `{"type":"response.created","response":{"id":"response-1","status":"in_progress"}}` + const delta = `{"type":"response.output_text.delta","delta":"héllo"}` + const completed = `{"type":"response.completed","response":{"id":"response-1","status":"completed","output":[{"type":"message","content":[{"type":"output_text","text":"héllo"}]}]}}` + for name, stream := range map[string]string{ + "streamed text": acpTestSSE(created, delta, completed, "[DONE]"), + "terminal fallback": acpTestSSE(completed), + } { + t.Run(name, func(t *testing.T) { + summary, err := acpParseFoundrySSE(strings.NewReader(stream)) + if err != nil || summary.Text != "héllo" || summary.ResponseID != "response-1" || summary.Status != "completed" { + t.Fatal("valid terminal response rejected") + } + }) + } + for name, stream := range map[string]string{ + "created only": acpTestSSE(created), + "partial text": acpTestSSE(created, delta), + "done without terminal": acpTestSSE(created, delta, "[DONE]"), + "created lies completed": acpTestSSE(`{"type":"response.created","response":{"id":"response-1","status":"completed"}}`), + "error after completed": acpTestSSE(created, delta, completed, `{"type":"error","error":{"message":"test-only-private-detail"}}`), + "duplicate terminal": acpTestSSE(completed, completed), + "missing terminal response": acpTestSSE(created, `{"type":"response.completed"}`), + "wrong terminal status": acpTestSSE(created, `{"type":"response.completed","response":{"id":"response-1","status":"in_progress"}}`), + "terminal error": acpTestSSE(`{"type":"response.completed","response":{"id":"response-1","status":"completed","error":{"message":"private"}}}`), + "terminal incomplete": acpTestSSE(`{"type":"response.completed","response":{"id":"response-1","status":"completed","incomplete_details":{"reason":"max_output_tokens"}}}`), + "wrong response identity": acpTestSSE(created, strings.Replace(completed, "response-1", "response-2", 1)), + "changed streamed text": acpTestSSE(created, strings.Replace(delta, "héllo", "wrong", 1), completed), + "native tool event": acpTestSSE(created, `{"type":"response.web_search_call.completed"}`, completed), + "native tool item": acpTestSSE(created, `{"type":"response.output_item.done","item":{"type":"web_search_call","id":"native"}}`, completed), + "duplicate JSON field": acpTestSSE(`{"type":"error","type":"response.completed","response":{"id":"response-1","status":"completed"}}`), + "truncated JSON": acpTestSSE(created, `{"type":"response.completed","response":`), + "truncated event": strings.TrimSuffix(acpTestSSE(completed), "\n"), + "missing delta": acpTestSSE(created, `{"type":"response.output_text.delta"}`, completed), + "null delta": acpTestSSE(created, `{"type":"response.output_text.delta","delta":null}`, completed), + "failed terminal": acpTestSSE(created, `{"type":"response.failed","response":{"id":"response-1","status":"failed"}}`), + "cancelled terminal": acpTestSSE(created, `{"type":"response.cancelled","response":{"id":"response-1","status":"cancelled"}}`), + "oversized output": acpTestSSE(created, `{"type":"response.output_text.delta","delta":"`+strings.Repeat("x", defaultMaxOutputBytes+1)+`"}`, completed), + } { + t.Run(name, func(t *testing.T) { + if _, err := acpParseFoundrySSE(strings.NewReader(stream)); err == nil { + t.Fatal("incomplete or malformed provider stream settled successfully") + } + }) + } +} + +func TestACPResponsesWaitForCompleteFunctionCallItems(t *testing.T) { + const added = `{"type":"response.output_item.added","item":{"id":"item-1","type":"function_call","name":"probe","call_id":"call-1","arguments":"","status":"in_progress"}}` + const delta = `{"type":"response.function_call_arguments.delta","item_id":"item-1","delta":"{"}` + const argsDone = `{"type":"response.function_call_arguments.done","item_id":"item-1","arguments":"{}"}` + const itemDone = `{"type":"response.output_item.done","item":{"id":"item-1","type":"function_call","name":"probe","call_id":"call-1","arguments":"{}","status":"completed"}}` + const completed = `{"type":"response.completed","response":{"id":"response-1","status":"completed"}}` + summary, err := acpParseFoundrySSE(strings.NewReader(acpTestSSE(added, delta, argsDone, itemDone, completed))) + if err != nil || len(summary.FunctionCalls) != 1 || summary.FunctionCalls[0].CallID != "call-1" { + t.Fatal("complete function call stream rejected") + } + for name, stream := range map[string]string{ + "only added": acpTestSSE(added, completed), + "partial arguments": acpTestSSE(added, delta, completed), + "arguments done without item": acpTestSSE(added, delta, argsDone, completed), + "item done without response terminal": acpTestSSE(itemDone), + "item still incomplete": acpTestSSE(strings.Replace(itemDone, `"status":"completed"`, `"status":"in_progress"`, 1), completed), + "terminal omits pending item": acpTestSSE(added, strings.Replace(itemDone, "item-1", "item-2", 1), completed), + } { + t.Run(name, func(t *testing.T) { + if _, err := acpParseFoundrySSE(strings.NewReader(stream)); err == nil { + t.Fatal("partial tool-call stream accepted") + } + }) + } +} + +func TestACPResponsesEventAndStreamBounds(t *testing.T) { + for name, stream := range map[string]string{ + "event count": strings.Repeat(acpTestSSE(`{"type":"response.reasoning_text.delta","delta":"x"}`), defaultMaxEvents+1), + "single event": acpTestSSE(`{"type":"response.reasoning_text.delta","delta":"` + strings.Repeat("x", defaultMaxEventBytes) + `"}`), + "stream bytes": strings.Repeat(":"+strings.Repeat("x", 1<<20)+"\n", 17), + } { + t.Run(name, func(t *testing.T) { + if _, err := acpParseFoundrySSE(strings.NewReader(stream)); err == nil { + t.Fatal("unbounded provider stream accepted") + } + }) + } +} + +func TestACPInvalidProviderCallsNeverReachMCP(t *testing.T) { + valid := acpTestCall("probe", "call-1", `{}`) + for name, calls := range map[string][]foundryOutputItem{ + "unknown tool": {acpTestCall("forbidden", "call-1", `{}`)}, + "duplicate call": {valid, valid}, + "missing call ID": {acpTestCall("probe", "", `{}`)}, + "array arguments": {acpTestCall("probe", "call-1", `[]`)}, + "null arguments": {acpTestCall("probe", "call-1", `null`)}, + "malformed arguments": {acpTestCall("probe", "call-1", `{`)}, + "duplicate arguments": {acpTestCall("probe", "call-1", `{"key":1,"key":2}`)}, + "unpaired surrogate": {acpTestCall("probe", "call-1", `{"key":"\ud800"}`)}, + "missing arguments": {{Type: "function_call", Name: "probe", CallID: "call-1"}}, + "bad sibling": {valid, acpTestCall("forbidden", "call-2", `{}`)}, + "native tool": {{Type: "web_search_call", ID: "native"}}, + } { + t.Run(name, func(t *testing.T) { + var requests atomic.Int32 + mcp := &acpTestMCP{ + tools: func() []map[string]any { return acpTestTools("probe") }, + execute: func(w http.ResponseWriter, _ *http.Request, id json.RawMessage, _ string, _ json.RawMessage) { + acpTestToolResult(w, id, "unexpected", false) + }, + } + peer := newACPTestPeer(t, toolSchemaModeProviderStatic, func(w http.ResponseWriter, r *http.Request) { + acpTestReadProvider(t, r) + requests.Add(1) + acpTestCompleted(w, "invalid-calls", "", calls...) + }, mcp) + acpAssertFailure(t, peer.reply(peer.prompt("invalid call"))) + if mcp.calls.Load() != 0 || requests.Load() != 1 || len(peer.events) != 0 { + t.Fatal("malformed batch admitted a tool or model replay") + } + }) + } +} + +func TestACPTruncatedStreamNeverExecutesCompleteToolItem(t *testing.T) { + var requests atomic.Int32 + mcp := &acpTestMCP{ + tools: func() []map[string]any { return acpTestTools("probe") }, + execute: func(w http.ResponseWriter, _ *http.Request, id json.RawMessage, _ string, _ json.RawMessage) { + acpTestToolResult(w, id, "unexpected", false) + }, + } + peer := newACPTestPeer(t, toolSchemaModeProviderStatic, func(w http.ResponseWriter, r *http.Request) { + acpTestReadProvider(t, r) + requests.Add(1) + w.Header().Set("Content-Type", "text/event-stream") + _, _ = fmt.Fprint(w, acpTestSSE(`{"type":"response.output_item.done","item":{"type":"function_call","name":"probe","call_id":"call-1","arguments":"{}"}}`)) + }, mcp) + acpAssertFailure(t, peer.reply(peer.prompt("truncated"))) + if requests.Load() != 1 || mcp.calls.Load() != 0 || len(peer.events) != 0 { + t.Fatal("truncated response admitted a tool call") + } +} + +func TestACPResponsesRejectAmbiguousFoldedEventFields(t *testing.T) { + const response = `{"id":"response-1","status":"completed","output":[{"type":"message","role":"assistant","content":[{"type":"output_text","text":"checked"}]}]}` + const changed = `{"id":"response-2","status":"completed","output":[{"type":"message","role":"user","content":[{"type":"output_text","text":"unchecked"}]}]}` + const completed = `{"type":"response.completed","response":{"id":"response-1","status":"completed"}}` + const call = `{"id":"item-1","type":"function_call","name":"probe","call_id":"call-1","arguments":"{}","status":"in_progress"}` + for name, stream := range map[string]string{ + "response identity and role": acpTestSSE(`{"type":"response.completed","response":` + response + `,"Response":` + changed + `}`), + "response reverse order": acpTestSSE(`{"type":"response.completed","Response":` + changed + `,"response":` + response + `}`), + "Unicode folded response": acpTestSSE(`{"type":"response.completed","response":` + response + `,"reſponſe":` + changed + `}`), + "escaped folded response": acpTestSSE(`{"type":"response.completed","response":` + response + `,"re\u017fpon\u017fe":` + changed + `}`), + "item skips status validation": acpTestSSE(`{"type":"response.output_item.done","item":{"type":"message"},"Item":`+call+`}`, completed), + "item reverse order": acpTestSSE(`{"type":"response.output_item.done","Item":`+call+`,"item":{"type":"message"}}`, completed), + "delta changes value": acpTestSSE(`{"type":"response.output_text.delta","delta":"checked","Delta":"unchecked"}`, completed), + "type changes terminal": acpTestSSE(`{"type":"error","Type":"response.completed","response":` + response + `}`), + "item identity aliases": acpTestSSE(`{"type":"response.function_call_arguments.done","item_id":"item-1","Item_ID":"other","arguments":"{}"}`, `{"type":"response.output_item.done","item":`+strings.Replace(call, "in_progress", "completed", 1)+`}`, completed), + } { + t.Run(name, func(t *testing.T) { + summary, err := acpParseFoundrySSE(strings.NewReader(stream)) + if err == nil || summary.ResponseID != "" || summary.Text != "" || len(summary.FunctionCalls) != 0 { + t.Fatal("ambiguous provider event exposed a response or tool call") + } + }) + } +} + +func TestACPResponsesFoldedEventCannotExecuteTool(t *testing.T) { + var requests atomic.Int32 + mcp := &acpTestMCP{ + tools: func() []map[string]any { return acpTestTools("probe") }, + execute: func(w http.ResponseWriter, _ *http.Request, id json.RawMessage, _ string, _ json.RawMessage) { + acpTestToolResult(w, id, "unexpected", false) + }, + } + peer := newACPTestPeer(t, toolSchemaModeProviderStatic, func(w http.ResponseWriter, r *http.Request) { + acpTestReadProvider(t, r) + if requests.Add(1) > 1 { + acpTestCompleted(w, "response-2", "unexpected") + return + } + w.Header().Set("Content-Type", "text/event-stream") + _, _ = fmt.Fprint(w, acpTestSSE( + `{"type":"response.output_item.done","item":{"type":"message"},"Item":{"id":"item-1","type":"function_call","name":"probe","call_id":"call-1","arguments":"{}","status":"in_progress"}}`, + `{"type":"response.completed","response":{"id":"response-1","status":"completed"}}`)) + }, mcp) + reply := peer.reply(peer.prompt("folded event validation")) + if requests.Load() != 1 || mcp.calls.Load() != 0 || len(peer.events) != 0 { + t.Fatal("ambiguous provider item admitted a tool or another model request") + } + acpAssertFailure(t, reply) +} + +func TestACPResponsesMatchSingleFoldedEnvelopeFields(t *testing.T) { + const response = `{"id":"response-1","status":"completed","output":[{"type":"message","content":[{"type":"output_text","text":"checked"}]}]}` + for name, stream := range map[string]string{ + "response": acpTestSSE(`{"Type":"response.completed","Reſponſe":` + response + `}`), + "delta": acpTestSSE(`{"type":"response.output_text.delta","Delta":"checked"}`, `{"type":"response.completed","response":`+response+`}`), + } { + t.Run(name, func(t *testing.T) { + summary, err := acpParseFoundrySSE(strings.NewReader(stream)) + if err != nil || summary.ResponseID != "response-1" || summary.Text != "checked" || len(summary.FunctionCalls) != 0 { + t.Fatal("unambiguous folded event fields changed the response") + } + }) + } + stream := acpTestSSE( + `{"type":"response.function_call_arguments.done","Item_ID":"item-1","arguments":"{\"Key\":1,\"key\":2}"}`, + `{"type":"response.output_item.done","Item":{"id":"item-1","type":"function_call","name":"probe","call_id":"call-1","arguments":"{\"Key\":1,\"key\":2}","status":"completed"}}`, + `{"Type":"response.completed","Response":{"id":"response-1","status":"completed"}}`) + summary, err := acpParseFoundrySSE(strings.NewReader(stream)) + if err != nil || len(summary.FunctionCalls) != 1 { + t.Fatal("unambiguous folded item fields rejected a complete tool call") + } + var arguments string + if json.Unmarshal(summary.FunctionCalls[0].Arguments, &arguments) != nil || arguments != `{"Key":1,"key":2}` { + t.Fatal("case-sensitive tool argument keys changed") + } +} + +func TestACPResponsesRejectAmbiguousFoldedOutput(t *testing.T) { + const userMessage = `[{"type":"message","role":"user","content":[{"type":"output_text","text":"unchecked"}]}]` + const emptyAssistant = `[{"type":"message","role":"assistant"}]` + const incompleteCall = `[{"id":"item-1","type":"function_call","status":"in_progress","name":"probe","call_id":"call-1","arguments":"{}"}]` + const emptyCall = `[{"type":"function_call"}]` + for name, fields := range map[string]string{ + "message content preserved": `"output":` + userMessage + `,"Output":` + emptyAssistant, + "message reverse casing": `"Output":` + userMessage + `,"output":` + emptyAssistant, + "escaped output alias": `"output":` + userMessage + `,"\u004futput":` + emptyAssistant, + "incomplete call preserved": `"output":` + incompleteCall + `,"Output":` + emptyCall, + "call reverse casing": `"Output":` + incompleteCall + `,"output":` + emptyCall, + } { + t.Run(name, func(t *testing.T) { + response := `{"id":"response-1","status":"completed",` + fields + `}` + document, err := acpDecodeFoundryResponse([]byte(response)) + if err == nil || document.ID != "" || len(document.Output) != 0 { + t.Error("ambiguous response output exposed a decoded document") + } + summary, err := acpParseFoundrySSE(strings.NewReader(acpTestSSE(`{"type":"response.completed","response":` + response + `}`))) + if err == nil || summary.ResponseID != "" || summary.Text != "" || len(summary.FunctionCalls) != 0 { + t.Error("ambiguous response output exposed a response or tool call") + } + }) + } +} + +func TestACPResponsesFoldedOutputCannotExecuteTool(t *testing.T) { + const response = `{"id":"response-1","status":"completed","output":[{"id":"item-1","type":"function_call","status":"in_progress","name":"probe","call_id":"call-1","arguments":"{}"}],"Output":[{"type":"function_call"}]}` + for _, mediaType := range []string{"application/json", "text/event-stream"} { + t.Run(mediaType, func(t *testing.T) { + var requests atomic.Int32 + mcp := &acpTestMCP{ + tools: func() []map[string]any { return acpTestTools("probe") }, + execute: func(w http.ResponseWriter, _ *http.Request, id json.RawMessage, _ string, _ json.RawMessage) { + acpTestToolResult(w, id, "unexpected", false) + }, + } + peer := newACPTestPeer(t, toolSchemaModeProviderStatic, func(w http.ResponseWriter, r *http.Request) { + acpTestReadProvider(t, r) + if requests.Add(1) > 1 { + acpTestCompleted(w, "response-2", "unexpected") + return + } + w.Header().Set("Content-Type", mediaType) + if mediaType == "application/json" { + _, _ = fmt.Fprint(w, response) + } else { + _, _ = fmt.Fprint(w, acpTestSSE(`{"type":"response.completed","response":`+response+`}`)) + } + }, mcp) + reply := peer.reply(peer.prompt("folded output validation")) + if requests.Load() != 1 || mcp.calls.Load() != 0 || len(peer.events) != 0 { + t.Fatalf("ambiguous response output admitted effects: provider requests=%d, tool calls=%d, events=%d", requests.Load(), mcp.calls.Load(), len(peer.events)) + } + acpAssertFailure(t, reply) + }) + } +} + +func TestACPResponsesSingleFoldedOutputPreservesValidatedItems(t *testing.T) { + const response = `{"id":"response-1","status":"completed","OuTpUt":[{"type":"message","role":"assistant","status":"completed","content":[{"type":"output_text","text":"checked"}]},{"id":"item-1","type":"function_call","status":"completed","name":"probe","call_id":"call-1","arguments":"{\"Key\":1,\"key\":2}"}]}` + document, err := acpDecodeFoundryResponse([]byte(response)) + if err != nil || len(document.Output) != 2 || document.Output[0].Type != "message" || document.Output[1].Type != "function_call" { + t.Fatal("single folded output lost validated items") + } + summary, err := acpParseFoundrySSE(strings.NewReader(acpTestSSE(`{"type":"response.completed","response":` + response + `}`))) + if err != nil || summary.ResponseID != "response-1" || summary.Text != "checked" || len(summary.FunctionCalls) != 1 || summary.FunctionCalls[0].Name != "probe" || summary.FunctionCalls[0].CallID != "call-1" { + t.Fatal("single folded output changed text or function-call identity") + } + var arguments string + if json.Unmarshal(summary.FunctionCalls[0].Arguments, &arguments) != nil || arguments != `{"Key":1,"key":2}` { + t.Fatal("single folded output changed case-sensitive tool arguments") + } +} diff --git a/acp_run.go b/acp_run.go new file mode 100644 index 0000000..acd3e82 --- /dev/null +++ b/acp_run.go @@ -0,0 +1,159 @@ +package main + +import ( + "context" + "strings" + "sync" +) + +const acpMaxResponseRounds = 32 + +func (s *acpServer) runPrompt(ctx context.Context, session *acpSession, prompt string) (string, string, error) { + tools, err := session.mcp.tools(ctx) + if err != nil { + return "", "", err + } + allowed := make(map[string]bool, len(tools)) + for _, tool := range tools { + allowed[tool.Name] = true + } + request := foundryResponseRequest{Input: prompt, PreviousResponseID: session.previous} + if s.cfg.agent.ToolSchemaMode == toolSchemaModeRequest { + request.Tools = tools + } + seen := make(map[string]bool) + var text strings.Builder + for range acpMaxResponseRounds { + if err := ctx.Err(); err != nil { + return "", "", err + } + summary, err := acpCreateResponse(ctx, s.cfg, s.client, request) + if err != nil || text.Len()+len(summary.Text) > defaultMaxOutputBytes { + return "", "", errACPProvider + } + text.WriteString(summary.Text) + if len(summary.FunctionCalls) == 0 { + return summary.ResponseID, text.String(), nil + } + if len(seen)+len(summary.FunctionCalls) > defaultMaxBrokeredCalls { + return "", "", errACPProvider + } + // Validate the whole batch before admitting its first side effect. + calls := summary.FunctionCalls + for i, call := range calls { + if validateFoundryFunctionName(call.Name) != nil || !allowed[call.Name] || + validateFoundryCallID(call.CallID) != nil || validateProviderIdentifier("call", call.CallID) != nil || seen[call.CallID] || len(call.Arguments) == 0 { + return "", "", errACPProvider + } + arguments, err := normalizeFoundryToolArguments(call.Arguments) + var object map[string]any + if err != nil || len(arguments) > defaultMaxBrokeredBytes || acpDecode(arguments, &object, false) != nil || object == nil { + return "", "", errACPProvider + } + calls[i].Arguments = arguments + seen[call.CallID] = true + } + outputs, err := s.executeTools(ctx, session, calls) + if err != nil { + return "", "", err + } + request.Input = outputs + request.PreviousResponseID = summary.ResponseID + } + return "", "", errACPProvider +} + +func (s *acpServer) executeTools(ctx context.Context, session *acpSession, calls []foundryOutputItem) ([]foundryFunctionOutput, error) { + ids := make([]string, len(calls)) + started := 0 + for i, call := range calls { + if err := ctx.Err(); err != nil { + for j := range started { + _ = s.toolEvent(session.id, ids[j], calls[j].Name, "failed") + } + return nil, err + } + ids[i] = acpOpaqueID("tool-") + if err := s.toolEvent(session.id, ids[i], call.Name, "in_progress"); err != nil { + return nil, err + } + started++ + } + group, cancel := context.WithCancel(ctx) + defer cancel() + outputs := make([]foundryFunctionOutput, len(calls)) + var wg sync.WaitGroup + var mu sync.Mutex + var firstError error + var total int + // Orka's loopback MCP session admits two outstanding calls by default. + slots := make(chan struct{}, 2) + for i, call := range calls { + wg.Go(func() { + var output string + var isError bool + var err error + var acquired bool + select { + case slots <- struct{}{}: + acquired = true + if group.Err() == nil { + output, isError, err = session.mcp.execute(group, call.Name, call.Arguments) + } else { + err = group.Err() + } + case <-group.Done(): + err = group.Err() + } + mu.Lock() + total += len(output) + if total > defaultMaxBrokeredTurnBytes && err == nil { + err = errACPMCP + } + if err != nil { + if firstError == nil { + firstError = err + } + // Signal before writing events: stdout backpressure must not keep + // a sibling HTTP request alive after protocol authority is lost. + cancel() + } + mu.Unlock() + // Revoke the batch before a released slot can admit queued work. + if acquired { + <-slots + } + status := "completed" + if err != nil || isError { + status = "failed" + } + if eventErr := s.toolEvent(session.id, ids[i], call.Name, status); eventErr != nil { + mu.Lock() + if firstError == nil { + firstError = eventErr + } + cancel() + mu.Unlock() + } + outputs[i] = foundryFunctionOutput{Type: "function_call_output", CallID: call.CallID, Output: output} + }) + } + wg.Wait() + if firstError != nil { + return nil, firstError + } + if err := ctx.Err(); err != nil { + return nil, err + } + return outputs, nil +} + +func (s *acpServer) toolEvent(sessionID, id, name, status string) error { + update := "tool_call_update" + if status == "in_progress" { + update = "tool_call" + } + return s.update(sessionID, map[string]string{ + "sessionUpdate": update, "toolCallId": id, "title": name, "kind": "other", "status": status, + }) +} diff --git a/acp_test.go b/acp_test.go new file mode 100644 index 0000000..966f822 --- /dev/null +++ b/acp_test.go @@ -0,0 +1,320 @@ +package main + +import ( + "context" + "encoding/json" + "io" + "net/http" + "net/http/httptest" + "os" + "regexp" + "strings" + "sync" + "sync/atomic" + "testing" + "time" +) + +type acpTestPeer struct { + t *testing.T + in *io.PipeWriter + out *io.PipeReader + decoder *json.Decoder + done chan error + events []map[string]any + writeMu sync.Mutex + session string + id int +} + +type acpTestMCP struct { + tools func() []map[string]any + execute func(http.ResponseWriter, *http.Request, json.RawMessage, string, json.RawMessage) + lists atomic.Int32 + calls atomic.Int32 + badAuth atomic.Bool +} + +func (m *acpTestMCP) handler(w http.ResponseWriter, r *http.Request) { + if r.Method != http.MethodPost || r.URL.Path != "/mcp" || r.Header.Get("Authorization") != "Bearer test-only-mcp-token" || + r.Header.Get("MCP-Protocol-Version") != acpMCPVersion { + m.badAuth.Store(true) + w.WriteHeader(http.StatusUnauthorized) + return + } + var request acpRequest + if json.NewDecoder(r.Body).Decode(&request) != nil { + w.WriteHeader(http.StatusBadRequest) + return + } + switch request.Method { + case "initialize": + acpTestMCPResult(w, request.ID, map[string]any{ + "protocolVersion": acpMCPVersion, "capabilities": map[string]any{"tools": map[string]any{}}, + }) + case "notifications/initialized": + w.WriteHeader(http.StatusAccepted) + case "tools/list": + m.lists.Add(1) + tools := []map[string]any{} + if m.tools != nil { + tools = m.tools() + } + acpTestMCPResult(w, request.ID, map[string]any{"tools": tools}) + case "tools/call": + var params struct { + Name string `json:"name"` + Args json.RawMessage `json:"arguments"` + } + if json.Unmarshal(request.Params, ¶ms) != nil || m.execute == nil { + w.WriteHeader(http.StatusBadRequest) + return + } + m.calls.Add(1) + m.execute(w, r, request.ID, params.Name, params.Args) + default: + w.WriteHeader(http.StatusBadRequest) + } +} + +func acpTestMCPResult(w http.ResponseWriter, id json.RawMessage, result any) { + w.Header().Set("Content-Type", "application/json") + _ = json.NewEncoder(w).Encode(map[string]any{"jsonrpc": "2.0", "id": id, "result": result}) +} + +func acpTestToolResult(w http.ResponseWriter, id json.RawMessage, text string, isError bool) { + acpTestMCPResult(w, id, map[string]any{"content": []map[string]string{{"type": "text", "text": text}}, "isError": isError}) +} + +func acpTestTools(names ...string) []map[string]any { + tools := make([]map[string]any, 0, len(names)) + for _, name := range names { + tools = append(tools, map[string]any{"name": name, "description": "Test tool", "inputSchema": map[string]any{"type": "object", "additionalProperties": true}}) + } + return tools +} + +func newACPTestPeer(t *testing.T, mode string, provider http.HandlerFunc, mcp *acpTestMCP) *acpTestPeer { + t.Helper() + providerServer := httptest.NewServer(provider) + mcpServer := httptest.NewServer(http.HandlerFunc(mcp.handler)) + data := acpTestConfigBytes(mode) + env := acpTestEnvironment(data) + env[acpProviderBaseEnv] = providerServer.URL + "/v1" + cfg, err := verifyACPConfiguration(data, func(key string) string { return env[key] }) + if err != nil { + t.Fatal(err) + } + in, input := io.Pipe() + output, out := io.Pipe() + p := &acpTestPeer{t: t, in: input, out: output, decoder: json.NewDecoder(output), done: make(chan error, 1)} + go func() { p.done <- serveACP(context.Background(), cfg, in, out) }() + t.Cleanup(func() { + _ = p.in.Close() + _ = p.out.Close() + select { + case <-p.done: + case <-time.After(5 * time.Second): + t.Error("ACP child did not join after transport close") + } + providerServer.Close() + mcpServer.Close() + if mcp.badAuth.Load() { + t.Error("MCP request lost scoped authentication or protocol metadata") + } + }) + response := p.call("initialize", map[string]any{"protocolVersion": 1, "clientCapabilities": map[string]any{}}) + result, ok := response["result"].(map[string]any) + if !ok || result["protocolVersion"] != float64(1) { + t.Fatal("ACP initialize failed") + } + caps := result["agentCapabilities"].(map[string]any) + if caps["loadSession"] != false || caps["mcpCapabilities"].(map[string]any)["http"] != true || + caps["promptCapabilities"].(map[string]any)["image"] != false { + t.Fatal("ACP capabilities claim unsupported functionality") + } + cwd, err := os.Getwd() + if err != nil { + t.Fatal(err) + } + response = p.call("session/new", map[string]any{ + "cwd": cwd, + "mcpServers": []map[string]any{{"type": "http", "name": "broker", "url": mcpServer.URL + "/mcp", + "headers": []map[string]string{{"name": "Authorization", "value": "Bearer test-only-mcp-token"}}}}, + }) + result, ok = response["result"].(map[string]any) + if !ok { + t.Fatal("ACP session/new failed") + } + p.session, ok = result["sessionId"].(string) + if !ok || !strings.HasPrefix(p.session, "foundry-") { + t.Fatal("ACP session identity missing") + } + return p +} + +func (p *acpTestPeer) send(message any) { + p.t.Helper() + p.writeMu.Lock() + defer p.writeMu.Unlock() + if err := json.NewEncoder(p.in).Encode(message); err != nil { + p.t.Fatal("ACP request write failed") + } +} + +func (p *acpTestPeer) read() map[string]any { + p.t.Helper() + type outcome struct { + value map[string]any + err error + } + done := make(chan outcome, 1) + go func() { + var value map[string]any + err := p.decoder.Decode(&value) + done <- outcome{value, err} + }() + select { + case result := <-done: + if result.err != nil { + p.t.Fatal("ACP response read failed") + } + if result.value["jsonrpc"] != "2.0" { + p.t.Fatal("invalid ACP response version") + } + if result.value["method"] == "session/update" { + params := result.value["params"].(map[string]any) + if params["sessionId"] != p.session { + p.t.Fatal("event belongs to another session") + } + p.events = append(p.events, params["update"].(map[string]any)) + } + return result.value + case <-time.After(5 * time.Second): + p.t.Fatal("ACP response did not settle") + return nil + } +} + +func (p *acpTestPeer) start(method string, params any) int { + p.id++ + p.send(map[string]any{"jsonrpc": "2.0", "id": p.id, "method": method, "params": params}) + return p.id +} + +func (p *acpTestPeer) reply(id int) map[string]any { + p.t.Helper() + for { + message := p.read() + if message["id"] == float64(id) { + return message + } + if message["method"] != "session/update" { + p.t.Fatal("unexpected ACP response identity") + } + } +} + +func (p *acpTestPeer) call(method string, params any) map[string]any { + p.t.Helper() + return p.reply(p.start(method, params)) +} + +func (p *acpTestPeer) prompt(text string) int { + return p.start("session/prompt", map[string]any{"sessionId": p.session, "prompt": []map[string]string{{"type": "text", "text": text}}}) +} + +func acpTestCompleted(w http.ResponseWriter, id, text string, calls ...foundryOutputItem) { + output := append([]foundryOutputItem(nil), calls...) + if text != "" { + output = append(output, foundryOutputItem{Type: "message", Content: []foundryOutputContent{{Type: "output_text", Text: text}}}) + } + w.Header().Set("Content-Type", "application/json") + _ = json.NewEncoder(w).Encode(foundryResponse{ID: id, Status: "completed", Output: output}) +} + +func acpTestReadProvider(t *testing.T, r *http.Request) map[string]json.RawMessage { + t.Helper() + if r.Method != http.MethodPost || r.URL.Path != "/v1/responses" || r.Header.Get("Authorization") != "Bearer test-only-proxy-token" { + t.Error("provider request escaped scoped endpoint or authentication") + } + var body map[string]json.RawMessage + if json.NewDecoder(r.Body).Decode(&body) != nil { + t.Error("invalid provider request") + } + for _, field := range []string{"agent_session_id", "conversation", "background"} { + if _, present := body[field]; present { + t.Error("child supplied privileged provider state") + } + } + if string(body["model"]) != `"test-model"` || string(body["stream"]) != "true" || string(body["store"]) != "true" { + t.Error("provider request lost pinned model or Responses settings") + } + return body +} + +func acpAssertStop(t *testing.T, response map[string]any, stop string) { + t.Helper() + result, ok := response["result"].(map[string]any) + if !ok || result["stopReason"] != stop || response["error"] != nil { + t.Fatalf("expected ACP stop reason %s", stop) + } +} + +func acpAssertFailure(t *testing.T, response map[string]any) { + t.Helper() + err, ok := response["error"].(map[string]any) + if !ok || err["code"] != float64(-32603) || response["result"] != nil || err["message"] != acpInternalError.Message { + t.Fatal("expected generic fatal ACP prompt error") + } +} + +func acpAssertToolEvents(t *testing.T, events []map[string]any, completed, failed int) { + t.Helper() + pending := make(map[string]bool) + seen := make(map[string]bool) + idPattern := regexp.MustCompile(`^tool-[a-f0-9]{32}$`) + gotCompleted, gotFailed := 0, 0 + for _, event := range events { + kind := event["sessionUpdate"] + if kind != "tool_call" && kind != "tool_call_update" { + continue + } + id, _ := event["toolCallId"].(string) + if !idPattern.MatchString(id) || event["kind"] != "other" || len(event) != 5 { + t.Fatal("tool lifecycle metadata is invalid or carries payload fields") + } + if kind == "tool_call" { + if event["status"] != "in_progress" || seen[id] { + t.Fatal("duplicate or invalid tool start") + } + pending[id], seen[id] = true, true + continue + } + if !pending[id] { + t.Fatal("tool terminal has no unique matching start") + } + delete(pending, id) + switch event["status"] { + case "completed": + gotCompleted++ + case "failed": + gotFailed++ + default: + t.Fatal("invalid tool terminal status") + } + } + if len(pending) != 0 || gotCompleted != completed || gotFailed != failed { + t.Fatalf("tool lifecycle pairing differs: completed=%d failed=%d pending=%d", gotCompleted, gotFailed, len(pending)) + } +} + +func acpOutput(events []map[string]any) string { + var text strings.Builder + for _, event := range events { + if event["sessionUpdate"] == "agent_message_chunk" { + text.WriteString(event["content"].(map[string]any)["text"].(string)) + } + } + return text.String() +} diff --git a/acp_tool_queue_test.go b/acp_tool_queue_test.go new file mode 100644 index 0000000..4f4aa97 --- /dev/null +++ b/acp_tool_queue_test.go @@ -0,0 +1,68 @@ +package main + +import ( + "context" + "fmt" + "io" + "net/http" + "strings" + "sync/atomic" + "testing" + "testing/synctest" +) + +func TestACPFatalToolBatchDoesNotAdmitQueuedCalls(t *testing.T) { + synctest.Test(t, func(t *testing.T) { + ctx, cancel := context.WithCancel(context.Background()) + defer cancel() + started, release := make(chan struct{}), make(chan struct{}) + var dispatched atomic.Int32 + client := &http.Client{Transport: brokerFixtureTransport(func(r *http.Request) (*http.Response, error) { + if dispatched.Add(1) == 2 { + close(started) + } + <-release + return &http.Response{StatusCode: http.StatusServiceUnavailable, + Header: http.Header{"Content-Type": {"application/json"}}, + Body: io.NopCloser(strings.NewReader("{}")), Request: r}, nil + })} + s := &acpServer{ctx: ctx, cancel: cancel, writes: make(chan acpWrite, 32)} + writerDone := make(chan struct{}) + go func() { + defer close(writerDone) + for { + select { + case write := <-s.writes: + write.done <- nil + case <-ctx.Done(): + return + } + } + }() + session := &acpSession{id: "queue-test", mcp: &acpMCPClient{ + url: "http://127.0.0.1/mcp", client: client, + }} + calls := make([]foundryOutputItem, 32) + for i := range calls { + calls[i] = foundryOutputItem{Name: "probe", CallID: fmt.Sprintf("call-%d", i), Arguments: []byte("{}")} + } + done := make(chan error, 1) + go func() { + _, err := s.executeTools(ctx, session, calls) + done <- err + }() + <-started + // Both admitted requests are held while the other calls queue. Let + // the fatal responses complete only after that state is established. + synctest.Wait() + close(release) + if err := <-done; err == nil { + t.Fatal("fatal batch returned success") + } + if got := dispatched.Load(); got != 2 { + t.Fatalf("fatal batch dispatched %d calls, want only the two originally admitted calls", got) + } + cancel() + <-writerDone + }) +} diff --git a/adapter.go b/adapter.go index f584e46..ebd8771 100644 --- a/adapter.go +++ b/adapter.go @@ -281,7 +281,7 @@ func (a *adapter) startTurn(request harness.StartTurnRequest) (*turnState, strin Input: request.Input.Prompt, PreviousResponseID: turn.previousResponse, AgentSessionID: turn.agentSessionID, - Tools: foundryToolSchemas(request), + Tools: a.providerToolSchemas(request), }) return turn, eventsPath, nil } @@ -740,7 +740,7 @@ func (a *adapter) prepareContinuationLocked(turn *turnState) (foundryResponseReq Input: outputs, PreviousResponseID: previousResponseID, AgentSessionID: turn.agentSessionID, - Tools: foundryToolSchemas(turn.request), + Tools: a.providerToolSchemas(turn.request), }, nil } @@ -1376,6 +1376,13 @@ func foundryToolSchemas(request harness.StartTurnRequest) []foundryToolSchema { return tools } +func (a *adapter) providerToolSchemas(request harness.StartTurnRequest) []foundryToolSchema { + if strings.EqualFold(strings.TrimSpace(a.cfg.toolSchemaMode), toolSchemaModeProviderStatic) { + return nil + } + return foundryToolSchemas(request) +} + func normalizeFoundryToolArguments(raw json.RawMessage) (json.RawMessage, error) { if len(raw) == 0 { return json.RawMessage(`{}`), nil diff --git a/adapter_test.go b/adapter_test.go index 1cd7d6a..6726ebe 100644 --- a/adapter_test.go +++ b/adapter_test.go @@ -182,6 +182,86 @@ func TestAdapterBrokeredFunctionCallContinuation(t *testing.T) { } } +func TestAdapterProviderStaticBrokeredFunctionCallContinuation(t *testing.T) { + backend := &fakeResponsesBackend{responses: []scriptedResponse{ + { + validate: func(request foundryResponseRequest) error { + if len(request.Tools) != 0 { + return errors.New("provider-static mode forwarded request tool schemas") + } + return nil + }, + run: func(_ context.Context, callbacks responseCallbacks) (foundryStreamSummary, error) { + if err := callbacks.OnCreated(foundryResponse{ID: "resp-static", Status: "in_progress", AgentSessionID: "session-static"}); err != nil { + return foundryStreamSummary{}, err + } + call := foundryOutputItem{Type: "function_call", CallID: "call-static", Name: "lookup_ticket", Arguments: json.RawMessage(`{"ticket":"INC-1"}`)} + if err := callbacks.OnFunctionCall(call); err != nil { + return foundryStreamSummary{}, err + } + return foundryStreamSummary{ResponseID: "resp-static", AgentSessionID: "session-static", Status: "completed", FunctionCalls: []foundryOutputItem{call}}, nil + }, + }, + { + validate: func(request foundryResponseRequest) error { + if request.PreviousResponseID != "resp-static" || request.AgentSessionID != "session-static" { + return errors.New("provider-static continuation identifiers missing") + } + if len(request.Tools) != 0 { + return errors.New("provider-static continuation forwarded request tool schemas") + } + outputs, ok := request.Input.([]foundryFunctionOutput) + if !ok || len(outputs) != 1 || outputs[0].CallID != "call-static" || outputs[0].Type != "function_call_output" { + return errors.New("provider-static function_call_output missing") + } + return nil + }, + run: textResponseScript("resp-static-final", "session-static", "ticket is open").run, + }, + }} + cfg := testConfig("http://127.0.0.1") + cfg.toolSchemaMode = toolSchemaModeProviderStatic + adapter := newAdapter(cfg, backend) + request := startRequest("provider-static", "provider-static-session") + request.ToolExecutionMode = harness.ToolExecutionModeBrokered + request.Input.Tools = []harness.ToolDefinition{{ + Name: "lookup_ticket", + Description: "Look up a ticket", + BrokeredClass: harness.BrokeredToolClassRead, + Parameters: json.RawMessage(`{"type":"object","properties":{"ticket":{"type":"string"}},"required":["ticket"]}`), + }} + turn, _, err := adapter.startTurn(request) + if err != nil { + t.Fatalf("start: %v", err) + } + waitForFrame(t, adapter, turn, harness.FrameToolCallRequested) + continueRequest := harness.ContinueTurnRequest{ + Version: harness.ProtocolVersion, + Namespace: request.Namespace, + TaskName: request.TaskName, + SessionName: request.SessionName, + RuntimeSessionID: request.RuntimeSessionID, + TurnID: request.TurnID, + CorrelationID: request.CorrelationID, + ToolResults: []harness.ToolCallResult{{ + Version: harness.ProtocolVersion, + RuntimeSessionID: request.RuntimeSessionID, + TurnID: request.TurnID, + ToolCallID: orkaToolCallID(request.RuntimeSessionID, request.TurnID, "call-static"), + IdempotencyKey: harness.ToolRequestIdempotencyKey(request.RuntimeSessionID, request.TurnID, orkaToolCallID(request.RuntimeSessionID, request.TurnID, "call-static")), + Approved: true, + Output: json.RawMessage(`{"status":"open"}`), + }}, + } + if err := adapter.continueTurn(continueRequest); err != nil { + t.Fatalf("continue: %v", err) + } + waitTurnDone(t, turn) + if got := completedResult(turn.frames); got != "ticket is open" { + t.Fatalf("result = %q", got) + } +} + func TestAdapterVersionPinCreatesSessionOnce(t *testing.T) { backend := &fakeResponsesBackend{createSessionID: "pinned-session", responses: []scriptedResponse{ { diff --git a/broker.go b/broker.go new file mode 100644 index 0000000..ce63798 --- /dev/null +++ b/broker.go @@ -0,0 +1,634 @@ +package main + +import ( + "bytes" + "context" + "crypto/subtle" + "encoding/json" + "errors" + "io" + "net/http" + "sync" + "time" +) + +// Preserve cleanup receipts even for owners already at the ordinary operation +// limit. Settlement cannot consume the final slot reserved for retirement. +const ( + brokerOperationLimit = 16384 + brokerSettlementOperationLimit = brokerOperationLimit + 1 + brokerRetirementOperationLimit = brokerSettlementOperationLimit + 1 +) + +type brokerActive struct { + prompt string + cancel context.CancelFunc +} + +type lifecycleBroker struct { + cfg brokerConfiguration + tokenProvider foundryTokenProvider + httpClient *http.Client + store *brokerStore + mu sync.Mutex + ledger *brokerLedger + storageError error + active map[string]brokerActive + workers map[string]bool + nextReconcile map[string]time.Time + ctx context.Context + cancel context.CancelFunc + wg sync.WaitGroup +} + +func newLifecycleBroker(ctx context.Context, cfg brokerConfiguration, provider foundryTokenProvider, client *http.Client) (*lifecycleBroker, error) { + if provider == nil || len(cfg.bearer) < 32 || !brokerDigestValid(cfg.configDigest) { + return nil, errBrokerInvalid + } + store, ledger, err := openBrokerStore(cfg.stateDir, cfg.configDigest) + if err != nil { + return nil, err + } + if cfg.operationTimeout <= 0 { + cfg.operationTimeout = 45 * time.Second + } + if client == nil { + client = newBrokerHTTPClient() + } else { + copy := *client + copy.CheckRedirect = func(*http.Request, []*http.Request) error { return errBrokerRemote } + client = © + } + ctx, cancel := context.WithCancel(ctx) + b := &lifecycleBroker{cfg: cfg, tokenProvider: provider, httpClient: client, store: store, ledger: ledger, + active: map[string]brokerActive{}, workers: map[string]bool{}, nextReconcile: map[string]time.Time{}, ctx: ctx, cancel: cancel} + // A restarted broker never resumes an old request. Recovery only closes its + // exact durable authority and obtains provider cleanup evidence. + err = b.commitLocked(func(next *brokerLedger) error { + for _, session := range next.Sessions { + for _, prompt := range session.Prompts { + if !prompt.Settled { + prompt.Closing = true + // Keep intent bytes unchanged: a valid older ledger may + // have no room for a longer state name. Without its original + // active request, intent is already ambiguous ownership. + } + } + } + return nil + }) + if err != nil { + cancel() + store.close() + return nil, err + } + b.wg.Add(1) + go b.sweep() + return b, nil +} + +func (b *lifecycleBroker) close() { + b.cancel() + b.mu.Lock() + for _, active := range b.active { + active.cancel() + } + b.mu.Unlock() + b.wg.Wait() + b.store.close() +} + +func (b *lifecycleBroker) commitLocked(change func(*brokerLedger) error) error { + return b.commitCapacityLocked(false, change) +} + +// Reserved growth may spend already retained space for original acceptance and +// cleanup. Ordinary admissions, renewals and completion mappings must replenish +// every owner's reserve. Recovery of an older ledger is not rejected merely +// because it predates that admission rule. +func (b *lifecycleBroker) commitCapacityLocked(ordinary bool, change func(*brokerLedger) error) error { + if b.storageError != nil { + return b.storageError + } + data, err := json.Marshal(b.ledger) + var next brokerLedger + if err != nil || json.Unmarshal(data, &next) != nil { + return b.failStorageLocked() + } + if err := change(&next); err != nil { + return err + } + nextData, err := json.Marshal(&next) + if err != nil { + return b.failStorageLocked() + } + limit := brokerMaxLedgerBytes + if ordinary && !bytes.Equal(data, nextData) { + limit -= brokerLedgerReserveBytes(&next) + } + // A definite capacity refusal precedes all disk I/O. Keep the old ledger + // and its healthy writer available for acceptance and cleanup. + if len(nextData) > limit { + return errBrokerCapacity + } + if err := b.store.saveBytes(nextData); err != nil { + return b.failStorageLocked() + } + b.ledger = &next + return nil +} + +func (b *lifecycleBroker) failStorageLocked() error { + b.storageError = errBrokerStorage + // Preserve the last durable owner and stop live requests. Cancellation is + // containment only; failed storage cannot publish settlement or retirement. + // Creation and cleanup outlive prompt cancellation, but cannot outlive + // the broker's ability to retain their acknowledgements. + b.cancel() + for _, active := range b.active { + active.cancel() + } + return b.storageError +} + +func brokerEnsureSession(next *brokerLedger, c brokerContext) (*brokerSession, error) { + key := brokerJSONDigest(c.Owner) + if session := next.Sessions[key]; session != nil { + return session, nil + } + if len(next.Sessions) >= 4096 { + return nil, errBrokerStorage + } + session := &brokerSession{Owner: c.Owner, CreateState: "none", Prompts: map[string]*brokerPrompt{}, + Responses: map[string]brokerResponseID{}, Operations: map[string]string{}} + next.Sessions[key] = session + return session, nil +} + +func brokerEnsurePrompt(session *brokerSession, c brokerContext) (*brokerPrompt, error) { + key := c.promptKey() + if prompt := session.Prompts[key]; prompt != nil { + return prompt, nil + } + if session.Retiring || session.Retired { + return nil, errBrokerClosed + } + for _, prompt := range session.Prompts { + if prompt.Identity.PromptID == c.PromptID || (prompt.Identity.TaskUID == c.TaskUID && prompt.Identity.TaskAttempt == c.TaskAttempt) { + return nil, errBrokerConflict + } + } + if current := session.Prompts[session.CurrentPrompt]; current != nil && !current.Settled { + return nil, errBrokerConflict + } + if len(session.Prompts) >= 4096 { + return nil, errBrokerStorage + } + expires, err := time.Parse(time.RFC3339Nano, c.LeaseExpiresAt) + if err != nil { + return nil, errBrokerInvalid + } + prompt := &brokerPrompt{Identity: c, LeaseGeneration: c.LeaseGeneration, LeaseExpiresAt: expires, + Invocations: map[uint64]*brokerInvocation{}} + session.Prompts[key] = prompt + session.CurrentPrompt = key + return prompt, nil +} + +func brokerRecordOperation(session *brokerSession, path string, c brokerContext) (bool, error) { + digest := brokerOperationDigest(path, c) + if previous, ok := session.Operations[c.OperationID]; ok { + if previous != digest { + return true, errBrokerConflict + } + return true, nil + } + limit := brokerOperationLimit + switch path { + case brokerSettlePath: + limit = brokerSettlementOperationLimit + case brokerRetirePath: + limit = brokerRetirementOperationLimit + } + if len(session.Operations) >= limit { + return false, errBrokerStorage + } + session.Operations[c.OperationID] = digest + return false, nil +} + +func (b *lifecycleBroker) ServeHTTP(w http.ResponseWriter, r *http.Request) { + w.Header().Set("Cache-Control", "no-store") + if r.URL.Path == "/healthz" && r.Method == http.MethodGet { + b.mu.Lock() + healthy := b.storageError == nil && b.ctx.Err() == nil + b.mu.Unlock() + if !healthy { + w.WriteHeader(http.StatusServiceUnavailable) + return + } + w.WriteHeader(http.StatusOK) + return + } + values := r.Header.Values("Authorization") + if len(values) != 1 || subtle.ConstantTimeCompare([]byte(values[0]), []byte("Bearer "+b.cfg.bearer)) != 1 { + w.WriteHeader(http.StatusUnauthorized) + return + } + validPath := r.URL.Path == brokerResponsesPath || r.URL.Path == brokerRenewPath || r.URL.Path == brokerSettlePath || + r.URL.Path == brokerRetirePath || r.URL.Path == brokerStatusPath + if !validPath || r.URL.RawQuery != "" || r.URL.RawPath != "" || r.Header.Get("Content-Encoding") != "" || + (r.URL.Path == brokerStatusPath && r.Method != http.MethodGet) || (r.URL.Path != brokerStatusPath && r.Method != http.MethodPost) { + w.WriteHeader(http.StatusBadRequest) + return + } + body, err := io.ReadAll(io.LimitReader(r.Body, maxFoundryPromptBytes+1)) + if err != nil || len(body) > maxFoundryPromptBytes { + w.WriteHeader(http.StatusBadRequest) + return + } + if (r.URL.Path == brokerStatusPath && len(body) != 0) || + (r.URL.Path != brokerResponsesPath && r.URL.Path != brokerStatusPath && string(body) != "{}") { + w.WriteHeader(http.StatusBadRequest) + return + } + c, contextDigest, err := brokerParseContext(r, body, b.cfg.configDigest, time.Now()) + if err != nil { + brokerWriteError(w, err) + return + } + if r.URL.Path == brokerResponsesPath { + b.serveResponses(w, r, c, body) + return + } + if r.URL.Path != brokerStatusPath { + if err := b.control(r.URL.Path, c); err != nil { + brokerWriteError(w, err) + return + } + } + if r.URL.Path == brokerSettlePath || r.URL.Path == brokerRetirePath { + b.startReconcile(brokerJSONDigest(c.Owner), true) + } + b.mu.Lock() + response := b.controlResponseLocked(c, contextDigest) + if r.URL.Path == brokerRenewPath && b.canAcknowledgeCreateRenewalLocked(c, response) { + response.State = "open" + } + storageErr := b.storageError + b.mu.Unlock() + if storageErr != nil { + brokerWriteError(w, storageErr) + return + } + status := http.StatusOK + if (r.URL.Path == brokerSettlePath && !response.SettlementProven) || (r.URL.Path == brokerRetirePath && !response.RetirementProven) { + status = http.StatusConflict + } + w.Header().Set("Content-Type", "application/json") + w.WriteHeader(status) + _ = json.NewEncoder(w).Encode(response) +} + +func brokerWriteError(w http.ResponseWriter, err error) { + status := http.StatusServiceUnavailable + switch { + case errors.Is(err, errBrokerInvalid): + status = http.StatusBadRequest + case errors.Is(err, errBrokerConflict), errors.Is(err, errBrokerPending), errors.Is(err, errBrokerAmbiguous): + status = http.StatusConflict + case errors.Is(err, errBrokerClosed): + status = http.StatusGone + } + w.Header().Set("Content-Type", "application/json") + w.WriteHeader(status) + _, _ = w.Write([]byte(`{"error":"foundry_broker_request_failed"}`)) +} + +func (b *lifecycleBroker) control(path string, c brokerContext) error { + b.mu.Lock() + defer b.mu.Unlock() + session := b.ledger.Sessions[brokerJSONDigest(c.Owner)] + ordinary := path == brokerRenewPath || session == nil || + (path == brokerSettlePath && session.Prompts[c.promptKey()] == nil) + err := b.commitCapacityLocked(ordinary, func(next *brokerLedger) error { + session, err := brokerEnsureSession(next, c) + if err != nil { + return err + } + duplicate, err := brokerRecordOperation(session, path, c) + if err != nil { + return err + } + if path == brokerRetirePath { + session.Retiring = true + for _, prompt := range session.Prompts { + prompt.Closing = true + } + return nil + } + if session.Retiring || session.Retired { + if path == brokerRenewPath { + return errBrokerClosed + } + } + prompt, err := brokerEnsurePrompt(session, c) + if err != nil { + return err + } + if path == brokerSettlePath { + prompt.Closing = true + return nil + } + if duplicate { + return nil + } + if prompt.Closing || prompt.Settled { + return errBrokerClosed + } + expires, _ := time.Parse(time.RFC3339Nano, c.LeaseExpiresAt) + // A renewal may establish an owner before inference. Its generation is + // already authoritative and need not start at one. + if prompt.Identity.OperationID == c.OperationID && prompt.LastSequence == 0 && prompt.LeaseGeneration == c.LeaseGeneration { + return nil + } + if c.LeaseGeneration != prompt.LeaseGeneration+1 || !prompt.LeaseExpiresAt.After(time.Now()) || expires.Before(prompt.LeaseExpiresAt) { + return errBrokerConflict + } + prompt.LeaseGeneration, prompt.LeaseExpiresAt = c.LeaseGeneration, expires + return nil + }) + if err == nil && (path == brokerSettlePath || path == brokerRetirePath) { + if active, ok := b.active[brokerJSONDigest(c.Owner)]; ok && (path == brokerRetirePath || active.prompt == c.promptKey()) { + active.cancel() + } + } + return err +} + +func (b *lifecycleBroker) controlResponseLocked(c brokerContext, contextDigest string) brokerControlResponse { + response := brokerControlResponse{Protocol: brokerProtocol, OwnerDigest: brokerJSONDigest(c.Owner), OperationID: c.OperationID, + ContextSHA256: contextDigest, State: "open"} + session := b.ledger.Sessions[response.OwnerDigest] + if session == nil { + return response + } + response.CreatePending = session.CreateState == "intent" + response.RemoteSessionCreated = session.CreateState == "known" || session.CreateState == "deleted" + response.RetirementProven = session.Retired + if session.Retiring { + response.State = "retiring" + } + if session.Retired { + response.State = "retired" + response.SettlementProven = true + response.ProofDigest = session.ProofDigest + } + for key, prompt := range session.Prompts { + if c.PromptID != "" && key != c.promptKey() { + continue + } + if c.PromptID != "" { + response.LeaseGeneration = prompt.LeaseGeneration + response.LeaseExpiresAt = prompt.LeaseExpiresAt.UTC().Format(time.RFC3339Nano) + response.SettlementProven = prompt.Settled + if prompt.Closing { + response.State = "settling" + } + if prompt.Settled { + response.State = "settled" + response.ProofDigest = prompt.ProofDigest + } + } + for _, invocation := range prompt.Invocations { + switch invocation.State { + case "uncertain": + response.AmbiguousInvocations++ + case "intent": + active, running := b.active[response.OwnerDigest] + if running && active.prompt == key && invocation.Sequence == prompt.LastSequence { + response.ActiveInvocations++ + } else { + response.AmbiguousInvocations++ + } + case "accepted", "reserved": + if !prompt.Settled { + response.ActiveInvocations++ + } + } + } + } + if response.CreatePending || response.AmbiguousInvocations > 0 { + response.State = "blocked" + response.SettlementProven = false + response.RetirementProven = false + response.ProofDigest = "" + } + return response +} + +// A live original create may outlast a prompt's first lease. Acknowledge only +// the exact renewed lease while that request still owns the reserved invocation. +// Creation remains pending, and status, settlement and retirement remain blocked. +func (b *lifecycleBroker) canAcknowledgeCreateRenewalLocked(c brokerContext, response brokerControlResponse) bool { + if b.ctx.Err() != nil || b.storageError != nil || response.State != "blocked" || !response.CreatePending || + response.AmbiguousInvocations != 0 || response.ActiveInvocations != 1 || response.LeaseGeneration != c.LeaseGeneration { + return false + } + key, promptKey := brokerJSONDigest(c.Owner), c.promptKey() + session := b.ledger.Sessions[key] + if session == nil || session.CreateState != "intent" || session.Retiring || session.Retired || session.CurrentPrompt != promptKey { + return false + } + prompt := session.Prompts[promptKey] + expires, err := time.Parse(time.RFC3339Nano, c.LeaseExpiresAt) + if err != nil || prompt == nil || prompt.Closing || prompt.Settled || prompt.LeaseGeneration != c.LeaseGeneration || + !prompt.LeaseExpiresAt.Equal(expires) || !prompt.LeaseExpiresAt.After(time.Now()) { + return false + } + active, ok := b.active[key] + if !ok || active.prompt != promptKey || len(prompt.Invocations) != 1 { + return false + } + invocation := prompt.Invocations[prompt.LastSequence] + return invocation != nil && invocation.State == "reserved" +} + +func (b *lifecycleBroker) sweep() { + defer b.wg.Done() + ticker := time.NewTicker(100 * time.Millisecond) + defer ticker.Stop() + for { + select { + case <-b.ctx.Done(): + return + case <-ticker.C: + } + b.mu.Lock() + needsExpiry := false + for _, session := range b.ledger.Sessions { + for _, prompt := range session.Prompts { + if !prompt.Closing && !prompt.LeaseExpiresAt.After(time.Now()) { + needsExpiry = true + } + } + } + if needsExpiry { + _ = b.commitLocked(func(next *brokerLedger) error { + for _, session := range next.Sessions { + for _, prompt := range session.Prompts { + if !prompt.LeaseExpiresAt.After(time.Now()) { + prompt.Closing = true + } + } + } + return nil + }) + } + keys := make([]string, 0, len(b.ledger.Sessions)) + for key, session := range b.ledger.Sessions { + needsCleanup := session.Retiring && !session.Retired + for promptKey, prompt := range session.Prompts { + if prompt.Closing && !prompt.Settled { + needsCleanup = true + if active, ok := b.active[key]; ok && active.prompt == promptKey { + active.cancel() + } + } + } + if needsCleanup { + keys = append(keys, key) + } + } + b.mu.Unlock() + for _, key := range keys { + b.startReconcile(key, false) + } + } +} + +func (b *lifecycleBroker) startReconcile(key string, force bool) { + b.mu.Lock() + if b.ctx.Err() != nil || b.storageError != nil || b.workers[key] || (!force && time.Now().Before(b.nextReconcile[key])) { + b.mu.Unlock() + return + } + b.workers[key] = true + b.wg.Add(1) + b.mu.Unlock() + go func() { + defer b.wg.Done() + ctx, cancel := context.WithTimeout(b.ctx, b.cfg.operationTimeout) + defer cancel() + b.reconcile(ctx, key) + b.mu.Lock() + delete(b.workers, key) + b.nextReconcile[key] = time.Now().Add(time.Second) + b.mu.Unlock() + }() +} + +func (b *lifecycleBroker) reconcile(ctx context.Context, key string) { + b.mu.Lock() + session := b.ledger.Sessions[key] + _, active := b.active[key] + if session == nil || session.Retired || active { + b.mu.Unlock() + return + } + // Capture exactly the closed prompts covered by this cleanup attempt. A + // repeated settlement of an older turn cannot settle a concurrently added + // turn using remote evidence obtained before that turn existed. + closing := make([]string, 0, len(session.Prompts)) + ambiguous := false + quiescent := true + for promptKey, prompt := range session.Prompts { + if prompt.Closing && !prompt.Settled { + closing = append(closing, promptKey) + for _, invocation := range prompt.Invocations { + switch invocation.State { + case "completed", "reserved", "rejected": + default: + quiescent = false + } + } + } + for _, invocation := range prompt.Invocations { + if invocation.State == "uncertain" || (invocation.State == "intent" && invocation.ResponseID == "") { + ambiguous = true + } + } + } + remoteID, createState, retiring := session.RemoteID, session.CreateState, session.Retiring + b.mu.Unlock() + if len(closing) == 0 && !retiring { + return + } + if createState == "intent" { + // A later GET cannot acknowledge the original CREATE or prove that it + // has finished. Keep its owner unresolved even if the session appears; + // only that original request's validated acknowledgement makes it known. + return + } + kind := "never-created" + if createState == "known" { + if retiring && !ambiguous { + // Deletion also terminates an acknowledged response. It remains + // retryable if a previous DELETE succeeded before proof persistence. + if b.remoteSessionDelete(ctx, remoteID) != nil { + return + } + kind = "remote-delete-204-get-404" + } else if quiescent && !ambiguous && len(closing) > 0 { + // Completed responses have durable validated links. Empty prompts + // and reserved/rejected invocations cannot add remote work once + // their active request is gone. Preserve compute and its history. + kind = "remote-responses-quiescent" + } else { + if b.remoteSessionStop(ctx, remoteID) != nil { + return + } + kind = "remote-stop-idle" + } + } + // Stop is useful containment, but does not prove an unacknowledged request + // cannot arrive later. Keep its durable owner and never issue deletion. + if ambiguous { + return + } + b.mu.Lock() + defer b.mu.Unlock() + _ = b.commitLocked(func(next *brokerLedger) error { + session := next.Sessions[key] + if _, active := b.active[key]; active { + return errBrokerPending + } + for _, promptKey := range closing { + prompt := session.Prompts[promptKey] + prompt.Settled = true + prompt.ProofDigest = brokerJSONDigest(struct { + Owner, Prompt, Remote, Kind, At string + }{key, promptKey, brokerSHA([]byte(remoteID)), kind, time.Now().UTC().Format(time.RFC3339Nano)}) + for _, invocation := range prompt.Invocations { + if invocation.State != "completed" { + invocation.State = "settled" + } + } + } + if !retiring { + return nil + } + for _, prompt := range session.Prompts { + if !prompt.Settled { + return errBrokerPending + } + } + if createState == "known" { + session.CreateState = "deleted" + } + session.Retired = true + session.ProofDigest = brokerJSONDigest(struct{ Owner, Remote, Kind, At string }{ + key, brokerSHA([]byte(remoteID)), kind, time.Now().UTC().Format(time.RFC3339Nano)}) + return nil + }) +} diff --git a/broker_byte_capacity_test.go b/broker_byte_capacity_test.go new file mode 100644 index 0000000..d8c707e --- /dev/null +++ b/broker_byte_capacity_test.go @@ -0,0 +1,397 @@ +package main + +import ( + "bytes" + "context" + "encoding/json" + "errors" + "fmt" + "net/http" + "os" + "path/filepath" + "strings" + "sync" + "testing" + "time" +) + +// Retained, valid synthetic history places the small public requests below at +// the byte boundary. The independent real-admission reproduction uses provider +// call maps; this fixture avoids replaying megabytes of historical responses in +// every capacity regression. It never removes an operation or an owner. +func brokerCapacityFillHistory(t *testing.T, b *lifecycleBroker, c brokerContext, target int, ordinary bool) { + t.Helper() + b.mu.Lock() + defer b.mu.Unlock() + err := b.commitCapacityLocked(ordinary, func(next *brokerLedger) error { + owner := c.Owner + owner.RuntimeSessionUID = "capacity-retained-history" + key := brokerJSONDigest(owner) + if next.Sessions[key] != nil { + return errBrokerConflict + } + history := &brokerSession{Owner: owner, CreateState: "none", Retiring: true, Retired: true, + ProofDigest: brokerSHA([]byte("synthetic-retained-proof")), Prompts: map[string]*brokerPrompt{}, + Responses: map[string]brokerResponseID{}, Operations: map[string]string{}} + next.Sessions[key] = history + data, err := json.Marshal(next) + if err != nil { + return err + } + size := len(data) + digest := brokerSHA([]byte("synthetic-retained-operation")) + for i := 0; i < brokerOperationLimit; i++ { + prefix := fmt.Sprintf("history-%05d-", i) + id := prefix + strings.Repeat("<", 512-len(prefix)) + encoded, _ := json.Marshal(id) + growth := len(encoded) + len(digest) + 4 // value quotes, colon, comma + if len(history.Operations) == 0 { + growth-- + } + if size+growth > target { + break + } + history.Operations[id] = digest + size += growth + } + // Fit the last entry with a mix of six-byte HTML escapes and ASCII. + // A few spare bytes are harmless; no identifier bound is weakened. + prefix := "capacity-tail-" + encodedLength := target - size - len(digest) - 6 + if len(history.Operations) == 0 { + encodedLength++ + } + for encodedLength >= len(prefix) { + rest := encodedLength - len(prefix) + id := prefix + strings.Repeat("<", rest/6) + strings.Repeat("x", rest%6) + if len(id) <= 512 { + history.Operations[id] = digest + break + } + encodedLength-- + } + data, err = json.Marshal(next) + if err != nil || !brokerLedgerValid(next, next.ConfigDigest) || len(data) > target || target-len(data) > 128 { + return errBrokerInvalid + } + return nil + }) + if err != nil { + t.Fatalf("could not establish the valid byte-boundary fixture: %v", err) + } +} + +func TestBrokerByteReserveCoversEscapedAcceptanceAndCleanup(t *testing.T) { + cfg := brokerConfiguration{configDigest: brokerSHA([]byte("capacity-bounds"))} + c := brokerTestContext(cfg) + c.BodySHA256 = brokerSHA([]byte("synthetic-body")) + c.OperationID = strings.Repeat("&", 512) + ledger := &brokerLedger{Version: 1, ConfigDigest: cfg.configDigest, Sessions: map[string]*brokerSession{}} + session, err := brokerEnsureSession(ledger, c) + if err != nil { + t.Fatal("could not establish bound owner") + } + if _, err := brokerRecordOperation(session, brokerResponsesPath, c); err != nil { + t.Fatal("could not establish bound operation") + } + prompt, err := brokerEnsurePrompt(session, c) + if err != nil { + t.Fatal("could not establish bound prompt") + } + prompt.LastSequence = c.InvocationSequence + invocation := &brokerInvocation{Sequence: c.InvocationSequence, OperationID: c.OperationID, + BodyDigest: c.BodySHA256, State: "intent"} + prompt.Invocations[c.InvocationSequence] = invocation + before, err := json.Marshal(ledger) + if err != nil || !brokerLedgerValid(ledger, cfg.configDigest) { + t.Fatal("initial bound ledger is invalid") + } + if got := brokerLedgerReserveBytes(ledger); got != brokerOwnerReserveBytes+brokerPrincipalReserveBytes { + t.Fatal("owner or first principal reserve is missing") + } + ledger.PrincipalDigest = brokerSHA([]byte("bound-principal")) + principal, _ := json.Marshal(ledger) + if len(principal)-len(before) > brokerPrincipalReserveBytes { + t.Fatal("first principal exceeds its reserve") + } + session.RemoteID, session.CreateState = brokerIdentityRemoteSession, "deleted" + session.Retiring, session.Retired = true, true + session.ProofDigest = brokerSHA([]byte("retirement-proof")) + prompt.Closing, prompt.Settled = true, true + prompt.ProofDigest = brokerSHA([]byte("settlement-proof")) + invocation.ResponseID = strings.Repeat("<", maxProviderIdentifierBytes) + invocation.ResponseAlias, invocation.State = "fr_11111111-1111-4111-8111-111111111111", "settled" + session.Responses[invocation.ResponseAlias] = brokerResponseID{RemoteID: invocation.ResponseID, PromptKey: c.promptKey()} + for _, path := range []string{brokerSettlePath, brokerRetirePath} { + control, body := brokerTestControlContext(path, c) + control.OperationID = strings.Repeat("<", 512) + if path == brokerRetirePath { + control.OperationID = strings.Repeat(">", 512) + } + control.BodySHA256 = brokerSHA(body) + if _, err := brokerRecordOperation(session, path, control); err != nil { + t.Fatal("maximum escaped cleanup operation was rejected") + } + } + after, err := json.Marshal(ledger) + if err != nil || !brokerLedgerValid(ledger, cfg.configDigest) { + t.Fatal("maximal acceptance and cleanup ledger is invalid") + } + // Final proof fields dominate the few bytes by which an intermediate state + // can be longer. Add an explicit 128-byte margin for those intermediate + // strings and the optional LastAlias that ordinary completion already sizes. + growth := len(after) - len(principal) + 128 + if growth > brokerOwnerReserveBytes { + t.Fatalf("bounded owner growth=%d exceeds reserve=%d", growth, brokerOwnerReserveBytes) + } + if brokerLedgerReserveBytes(ledger) != 0 { + t.Fatal("retired ownership retained an unnecessary future-growth reserve") + } + t.Logf("principalGrowthBytes=%d ownerGrowthWithMarginBytes=%d ownerReserveBytes=%d", len(principal)-len(before), growth, brokerOwnerReserveBytes) +} + +func TestBrokerByteCapacityRejectsBeforeIOAndKeepsFailureClosed(t *testing.T) { + b, c := newBrokerResponseIdentityFixture(t) + brokerCapacityFillHistory(t, b, c, brokerMaxLedgerBytes-brokerOwnerReserveBytes-512, true) + before := brokerIdentityLedgerBytes(t, b) + path := filepath.Join(b.store.dir, "state.json") + record, err := os.Open(path) + if err != nil { + t.Fatal("could not pin original ledger") + } + defer func() { _ = record.Close() }() + oldInfo, _ := record.Stat() + active, cancel := context.WithCancel(context.Background()) + defer cancel() + b.active[brokerJSONDigest(c.Owner)] = brokerActive{prompt: c.promptKey(), cancel: cancel} + // A write would fail here. Capacity rejection must happen first and must + // leave this original active request and the durable writer healthy. + if os.Rename(b.store.dir, b.store.dir+"-hidden") != nil { + t.Fatal("could not hide the fixture directory") + } + defer func() { _ = os.Rename(b.store.dir+"-hidden", b.store.dir) }() + renewal, body := brokerTestControlContext(brokerRenewPath, c) + renewal.OperationID = strings.Repeat("<", 512) + renewal.BodySHA256 = brokerSHA(body) + renewal.LeaseGeneration++ + if err := b.control(brokerRenewPath, renewal); !errors.Is(err, errBrokerCapacity) { + t.Fatalf("ordinary growth did not stop before I/O: %v", err) + } + err = b.commitLocked(func(next *brokerLedger) error { + session := next.Sessions[brokerJSONDigest(c.Owner)] + for i := range 32 { + id := fmt.Sprintf("extra-%02d-", i) + strings.Repeat("<", 500) + session.Operations[id] = brokerSHA([]byte("extra-operation")) + } + return nil + }) + if !errors.Is(err, errBrokerCapacity) || b.storageError != nil || active.Err() != nil || + brokerSHA(before) != brokerJSONDigest(b.ledger) { + t.Fatal("hard-cap refusal poisoned the writer, cancelled work, or mutated ownership") + } + if os.Rename(b.store.dir+"-hidden", b.store.dir) != nil { + t.Fatal("could not restore the fixture directory") + } + newInfo, err := os.Stat(path) + if err != nil || !os.SameFile(oldInfo, newInfo) || !bytes.Equal(before, brokerIdentityLedgerBytes(t, b)) { + t.Fatal("capacity rejection rewrote the ledger") + } + if os.Rename(b.store.dir, b.store.dir+"-hidden") != nil { + t.Fatal("could not inject a real storage failure") + } + err = b.commitLocked(func(next *brokerLedger) error { + next.Sessions[brokerJSONDigest(c.Owner)].Prompts[c.promptKey()].Closing = true + return nil + }) + if !errors.Is(err, errBrokerStorage) || !errors.Is(b.storageError, errBrokerStorage) || active.Err() != context.Canceled { + t.Fatal("actual I/O failure did not remain fail-closed") + } + if brokerSHA(before) != brokerJSONDigest(b.ledger) { + t.Fatal("failed I/O replaced original durable ownership") + } +} + +func brokerCapacityControlProof(t *testing.T, base, path string, c brokerContext) brokerControlResponse { + t.Helper() + // Race instrumentation must repeatedly encode and decode the full 32 MiB + // history. Give that work a bounded window on slower native builders. + deadline := time.Now().Add(90 * time.Second) + for { + status, data, err := brokerTestHTTP(context.Background(), base, path, c, []byte("{}")) + var proof brokerControlResponse + if err != nil || json.Unmarshal(data, &proof) != nil || (status != http.StatusOK && status != http.StatusConflict) { + t.Fatalf("bounded cleanup failed at byte capacity: status=%d", status) + } + if status == http.StatusOK { + return proof + } + if time.Now().After(deadline) { + t.Fatal("bounded cleanup did not reach proof") + } + time.Sleep(20 * time.Millisecond) + } +} + +func TestBrokerByteCapacityConcurrentOwnersPreserveAcceptanceAndCleanup(t *testing.T) { + responseID := strings.Repeat("<", maxProviderIdentifierBytes) + f := newBrokerEvidenceFixture(t, func(request foundryResponseRequest) (string, []byte) { + response := foundryResponse{ID: responseID, AgentSessionID: request.AgentSessionID, Status: "completed"} + for i := range defaultMaxBrokeredCalls { + prefix := fmt.Sprintf("call-%03d-", i) + response.Output = append(response.Output, foundryOutputItem{ID: fmt.Sprintf("item-%03d", i), Type: "function_call", + Name: "hosted-probe-read", CallID: prefix + strings.Repeat("x", maxProviderIdentifierBytes-len(prefix)), Arguments: json.RawMessage(`"{}"`)}) + } + data, err := json.Marshal(response) + if err != nil { + t.Error("could not encode bounded fixture response") + } + return "application/json", data + }) + started, release := make(chan struct{}, 2), make(chan struct{}) + var releaseOnce sync.Once + unblock := func() { releaseOnce.Do(func() { close(release) }) } + defer unblock() + f.createCheck = func(string) { started <- struct{}{}; <-release } + cfg := brokerTestConfig(t, f) + cfg.operationTimeout = 2 * time.Minute + b, server := startBrokerTest(t, cfg) + owners := make([]brokerContext, 2) + results := make([]<-chan brokerHTTPResult, len(owners)) + for i := range owners { + c := brokerTestContext(cfg) + c.Owner.RuntimeSessionUID = fmt.Sprintf("capacity-owner-%d", i) + c.OperationID = fmt.Sprintf("foundry-%064x", i+1) + c.LeaseExpiresAt = time.Now().Add(4 * time.Minute).UTC().Format(time.RFC3339Nano) + owners[i] = c + results[i] = brokerAsyncInference(context.Background(), server.URL, c, brokerTestBody("")) + } + for range owners { + select { + case <-started: + case <-time.After(10 * time.Second): + t.Fatal("original creates did not reach the acknowledgement gate") + } + } + b.mu.Lock() + reserve := brokerLedgerReserveBytes(b.ledger) + for _, c := range owners { + if b.ledger.Sessions[brokerJSONDigest(c.Owner)].CreateState != "intent" { + b.mu.Unlock() + t.Fatal("original creation intent was not durable before submission") + } + } + b.mu.Unlock() + if reserve != len(owners)*brokerOwnerReserveBytes { + t.Fatal("concurrent ownership was not fully reserved") + } + brokerCapacityFillHistory(t, b, owners[0], brokerMaxLedgerBytes-reserve-128, true) + unblock() + for _, result := range results { + select { + case response := <-result: + if response.err != nil || response.status != http.StatusServiceUnavailable { + t.Fatal("oversized completion mapping was exposed or original acceptance failed") + } + case <-time.After(2 * time.Minute): + t.Fatal("bounded inference did not return after capacity refusal") + } + } + b.mu.Lock() + valid, healthy := brokerLedgerValid(b.ledger, cfg.configDigest), b.storageError == nil + for _, c := range owners { + session := b.ledger.Sessions[brokerJSONDigest(c.Owner)] + invocation := session.Prompts[c.promptKey()].Invocations[c.InvocationSequence] + link := session.Responses[invocation.ResponseAlias] + valid = valid && session.CreateState == "known" && invocation.ResponseID == responseID && + link.RemoteID == responseID && !link.Completed && len(link.CallIDs) == 0 + } + b.mu.Unlock() + if !valid || !healthy { + t.Fatal("capacity withholding lost original acknowledgement or poisoned the writer") + } + proofs := make([]string, len(owners)) + for i, c := range owners { + settlement, _ := brokerTestControlContext(brokerSettlePath, c) + settlement.OperationID = strings.Repeat("<", 512) + for range 3 { + proof := brokerCapacityControlProof(t, server.URL, brokerSettlePath, settlement) + if !proof.SettlementProven || proof.ActiveInvocations != 0 || proof.AmbiguousInvocations != 0 { + t.Fatal("exact settlement retries consumed another owner's cleanup space") + } + } + retirement, _ := brokerTestControlContext(brokerRetirePath, c) + retirement.OperationID = strings.Repeat(">", 512) + proof := brokerCapacityControlProof(t, server.URL, brokerRetirePath, retirement) + if !proof.RetirementProven { + t.Fatal("owner could not persist exact retirement proof") + } + proofs[i] = proof.ProofDigest + } + before := brokerIdentityLedgerBytes(t, b) + b.close() + server.Close() + reopened, restarted := startBrokerTest(t, cfg) + for i, c := range owners { + retirement, _ := brokerTestControlContext(brokerRetirePath, c) + retirement.OperationID = strings.Repeat(">", 512) + proof := brokerCapacityControlProof(t, restarted.URL, brokerRetirePath, retirement) + if !proof.RetirementProven || proof.ProofDigest != proofs[i] { + t.Fatal("restart lost the original retirement receipt") + } + status, _, err := brokerTestHTTP(context.Background(), restarted.URL, brokerResponsesPath, c, brokerTestBody("")) + if err != nil || status != http.StatusGone { + t.Fatal("restart admitted inference on retired ownership") + } + } + if !bytes.Equal(before, brokerIdentityLedgerBytes(t, reopened)) { + t.Fatal("exact retry or rejected replay changed retained capacity history") + } + creates, inferences, stops, deletes := f.counts() + if creates != len(owners) || inferences != len(owners) || stops != len(owners) || deletes != len(owners) { + t.Fatal("capacity recovery replayed a create, inference or cleanup") + } + t.Logf("owners=%d reservedBytes=%d ledgerBytes=%d accepted=%d withheld=%d settled=%d retired=%d replayed=0", len(owners), reserve, len(before), inferences, inferences, stops, deletes) +} + +func TestBrokerByteCapacityLegacyLedgerStillRecovers(t *testing.T) { + f := newBrokerFixture(t, "success") + cfg := brokerTestConfig(t, f) + b, server := startBrokerTest(t, cfg) + c := brokerTestContext(cfg) + c.LeaseExpiresAt = time.Now().Add(4 * time.Minute).UTC().Format(time.RFC3339Nano) + _ = brokerTestControl(t, server.URL, brokerRenewPath, c) + _ = brokerTestControl(t, server.URL, brokerSettlePath, c) + // An old valid ledger need not contain the newly required admission reserve. + // Preserve it without erasure or newly fabricated ownership on reopen. + brokerCapacityFillHistory(t, b, c, brokerMaxLedgerBytes-8192, false) + before := brokerIdentityLedgerBytes(t, b) + b.close() + server.Close() + reopened, restarted := startBrokerTest(t, cfg) + if !bytes.Equal(before, brokerIdentityLedgerBytes(t, reopened)) { + t.Fatal("legacy recovery rewrote existing ownership") + } + status, _, err := brokerTestHTTP(context.Background(), restarted.URL, brokerResponsesPath, c, brokerTestBody("")) + if err != nil || status != http.StatusGone { + t.Fatal("legacy recovery replayed a closed prompt") + } + fresh := c + fresh.Owner.RuntimeSessionUID = "new-owner-after-legacy-capacity" + fresh.OperationID = "new-ordinary-operation" + status, _, err = brokerTestHTTP(context.Background(), restarted.URL, brokerResponsesPath, fresh, brokerTestBody("")) + if err != nil || status != http.StatusServiceUnavailable { + t.Fatal("legacy byte-capacity state admitted unreserved ownership") + } + retirement, _ := brokerTestControlContext(brokerRetirePath, c) + proof := brokerCapacityControlProof(t, restarted.URL, brokerRetirePath, retirement) + if !proof.RetirementProven { + t.Fatal("existing legacy cleanup was rejected by the new reserve rule") + } + reopened.mu.Lock() + healthy := reopened.storageError == nil && reopened.ledger.Sessions[brokerJSONDigest(fresh.Owner)] == nil + reopened.mu.Unlock() + creates, inferences, stops, deletes := f.counts() + if !healthy || creates+inferences+stops+deletes != 0 { + t.Fatal("legacy capacity refusal poisoned the writer or submitted remote work") + } +} diff --git a/broker_create_ack_test.go b/broker_create_ack_test.go new file mode 100644 index 0000000..798227e --- /dev/null +++ b/broker_create_ack_test.go @@ -0,0 +1,61 @@ +package main + +import ( + "context" + "encoding/json" + "net/http" + "testing" + "time" +) + +func TestBrokerLostCreateAckCannotRetireFromLaterGET(t *testing.T) { + f := newBrokerFixture(t, "hold-create") + defer f.unblock() + cfg := brokerTestConfig(t, f) + b, server := startBrokerTest(t, cfg) + c := brokerTestContext(cfg) + done := brokerAsyncInference(context.Background(), server.URL, c, brokerTestBody("")) + select { + case <-f.started: + case <-time.After(3 * time.Second): + t.Fatal("original CREATE did not reach its held acknowledgement") + } + result := brokerWaitInference(t, done) + if result.err != nil || result.status != http.StatusConflict { + t.Fatalf("original CREATE timeout HTTP = %d, want ambiguous 409", result.status) + } + brokerAwait(t, func() bool { + b.mu.Lock() + defer b.mu.Unlock() + return len(b.workers) == 0 + }) + retire, body := brokerTestControlContext(brokerRetirePath, c) + status, _, err := brokerTestHTTP(context.Background(), server.URL, brokerRetirePath, retire, body) + if err != nil || status != http.StatusOK && status != http.StatusConflict { + t.Fatal("original retirement request failed unexpectedly") + } + brokerAwait(t, func() bool { + b.mu.Lock() + defer b.mu.Unlock() + return len(b.workers) == 0 + }) + statusContext, body := brokerTestControlContext(brokerStatusPath, c) + status, data, err := brokerTestHTTP(context.Background(), server.URL, brokerStatusPath, statusContext, body) + var proof brokerControlResponse + if err != nil || status != http.StatusOK || json.Unmarshal(data, &proof) != nil { + t.Fatal("original owner status is unreadable") + } + select { + case <-f.release: + t.Fatal("original CREATE acknowledgement was unexpectedly released") + default: + } + creates, inference, stops, deletes := f.counts() + if creates != 1 || inference != 0 { + t.Fatal("reproducer replayed CREATE or sent inference") + } + t.Logf("originalCREATEHeld=true creates=%d inference=%d stops=%d deletes=%d createPending=%v retired=%v", creates, inference, stops, deletes, proof.CreatePending, proof.RetirementProven) + if !proof.CreatePending || proof.SettlementProven || proof.RetirementProven || proof.ProofDigest != "" || deletes != 0 { + t.Error("later GET authorized retirement before the original CREATE acknowledgement") + } +} diff --git a/broker_create_lease_test.go b/broker_create_lease_test.go new file mode 100644 index 0000000..c129cdd --- /dev/null +++ b/broker_create_lease_test.go @@ -0,0 +1,90 @@ +package main + +import ( + "context" + "errors" + "sync/atomic" + "testing" + "time" +) + +func TestBrokerCreationRequiresCurrentLeaseBeforeIntent(t *testing.T) { + for _, renewed := range []bool{false, true} { + name := "expired" + if renewed { + name = "renewed-during-preflight" + } + t.Run(name, func(t *testing.T) { + f := newBrokerFixture(t, "success") + cfg := brokerTestConfig(t, f) + c := brokerTestContext(cfg) + c.LeaseExpiresAt = time.Now().Add(-time.Second).UTC().Format(time.RFC3339Nano) + body := brokerTestBody("") + c.BodySHA256 = brokerSHA(body) + var request acpResponseRequest + if acpDecode(body, &request, true) != nil { + t.Fatal("invalid inference fixture") + } + store, ledger, err := openBrokerStore(cfg.stateDir, cfg.configDigest) + if err != nil { + t.Fatal("could not open lease fixture store") + } + ctx, cancel := context.WithCancel(t.Context()) + b := &lifecycleBroker{cfg: cfg, store: store, ledger: ledger, ctx: ctx, cancel: cancel, + httpClient: newBrokerHTTPClient(), active: map[string]brokerActive{}, + workers: map[string]bool{}, nextReconcile: map[string]time.Time{}} + t.Cleanup(b.close) + // Reproduce a previously admitted invocation whose lease expires before + // CREATE preflight finishes. No sweep runs in this fixture: admission + // must check the timestamp without waiting for background cleanup. + err = b.commitLocked(func(next *brokerLedger) error { + session, ensureErr := brokerEnsureSession(next, c) + if ensureErr != nil { + return ensureErr + } + if _, ensureErr = brokerRecordOperation(session, brokerResponsesPath, c); ensureErr != nil { + return ensureErr + } + prompt, ensureErr := brokerEnsurePrompt(session, c) + if ensureErr != nil { + return ensureErr + } + prompt.LastSequence = c.InvocationSequence + prompt.Invocations[c.InvocationSequence] = &brokerInvocation{Sequence: c.InvocationSequence, + OperationID: c.OperationID, BodyDigest: c.BodySHA256, State: "reserved"} + return nil + }) + if err != nil || !brokerLedgerValid(b.ledger, cfg.configDigest) { + t.Fatal("invalid admitted invocation fixture") + } + var tokens atomic.Int32 + b.tokenProvider = brokerEvidenceTokenProvider(func(context.Context) (string, error) { + if tokens.Add(1) == 3 && renewed { + b.mu.Lock() + err := b.commitLocked(func(next *brokerLedger) error { + prompt := next.Sessions[brokerJSONDigest(c.Owner)].Prompts[c.promptKey()] + prompt.LeaseGeneration++ + prompt.LeaseExpiresAt = time.Now().Add(10 * time.Second) + return nil + }) + b.mu.Unlock() + if err != nil { + return "", err + } + } + return brokerTestToken(), nil + }) + _, err = b.invoke(ctx, c, request.foundryResponseRequest) + creates, inferences, _, _ := f.counts() + owner := b.ledger.Sessions[brokerJSONDigest(c.Owner)] + if renewed { + if err != nil || creates != 1 || inferences != 1 || owner.Prompts[c.promptKey()].LeaseGeneration != 2 { + t.Fatal("current renewed lease did not authorize the original creation") + } + } else if !errors.Is(err, errBrokerClosed) || creates != 0 || inferences != 0 || + owner.CreateState != "none" || owner.RemoteID != "" { + t.Fatalf("expired lease reached creation: creates=%d inferences=%d state=%s", creates, inferences, owner.CreateState) + } + }) + } +} diff --git a/broker_evidence_test.go b/broker_evidence_test.go new file mode 100644 index 0000000..a317811 --- /dev/null +++ b/broker_evidence_test.go @@ -0,0 +1,250 @@ +package main + +import ( + "bytes" + "context" + "encoding/base64" + "encoding/json" + "errors" + "net/http" + "net/http/httptest" + "os" + "path/filepath" + "strings" + "sync/atomic" + "testing" +) + +func newBrokerEvidenceFixture(t *testing.T, reply func(foundryResponseRequest) (string, []byte)) *brokerFixture { + t.Helper() + f := newBrokerFixture(t, "success") + 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 + } + if r.Header.Get("Authorization") != "Bearer "+brokerTestToken() || r.URL.Query().Get("api-version") != "v1" { + t.Error("inference lacked the fixture's exact provider authority") + w.WriteHeader(http.StatusUnauthorized) + return + } + var request foundryResponseRequest + if json.NewDecoder(r.Body).Decode(&request) != nil { + t.Error("inference request was not decodable") + w.WriteHeader(http.StatusBadRequest) + return + } + f.mu.Lock() + _, exists := f.sessions[request.AgentSessionID] + if exists { + f.sessions[request.AgentSessionID] = "active" + f.inferences++ + f.requests = append(f.requests, request) + } + f.mu.Unlock() + if !exists { + t.Error("inference lacked an owned remote session") + w.WriteHeader(http.StatusNotFound) + return + } + media, body := reply(request) + w.Header().Set("Content-Type", media) + _, _ = w.Write(body) + })) + return f +} + +func TestBrokerMalformedSSECannotAcknowledgeInference(t *testing.T) { + for _, name := range []string{ + "missing-type", "unknown-type", "wrong-event", "missing-status", "unknown-status", + "created-completed", "queued-in-progress", "completed-in-progress", "completed-error", "wrong-session", + } { + t.Run(name, func(t *testing.T) { + f := newBrokerEvidenceFixture(t, func(request foundryResponseRequest) (string, []byte) { + response := map[string]any{"id": "provider-evidence", "status": "in_progress", "agent_session_id": request.AgentSessionID, "output": []any{}} + event := map[string]any{"type": "response.created", "response": response} + switch name { + case "missing-type": + delete(event, "type") + case "unknown-type": + event["type"] = "not-a-response-event" + case "wrong-event": + event["type"] = "response.output_text.delta" + event["delta"] = "fixture" + case "missing-status": + delete(response, "status") + case "unknown-status": + response["status"] = "not-a-response-status" + case "created-completed": + response["status"] = "completed" + case "queued-in-progress": + event["type"] = "response.queued" + case "completed-in-progress": + event["type"] = "response.completed" + case "completed-error": + event["type"], response["status"] = "response.completed", "completed" + response["error"] = map[string]string{"code": "server_error", "message": "fixture-error-do-not-persist"} + case "wrong-session": + response["agent_session_id"] = "another-session" + } + data, _ := json.Marshal(event) + return "text/event-stream", []byte("data: " + string(data) + "\n\n") + }) + cfg := brokerTestConfig(t, f) + b, server := startBrokerTest(t, cfg) + c := brokerTestContext(cfg) + status, _, err := brokerTestHTTP(context.Background(), server.URL, brokerResponsesPath, c, brokerTestBody("")) + if err != nil || status == http.StatusOK { + t.Fatal("malformed inference exposed a successful response") + } + if state := brokerInvocationState(b, c); state != "uncertain" { + t.Fatalf("malformed SSE fabricated acknowledgement: state=%s", state) + } + brokerPendingControl(t, server.URL, brokerSettlePath, c, false, 1) + brokerPendingControl(t, server.URL, brokerRetirePath, c, false, 1) + creates, inferences, _, deletes := f.counts() + if creates != 1 || inferences != 1 || deletes != 0 { + t.Fatal("malformed SSE replayed work or authorized deletion") + } + }) + } +} + +func TestBrokerFailureResponseRetainsAcknowledgement(t *testing.T) { + for _, media := range []string{"application/json", "text/event-stream"} { + for _, state := range []string{"failed", "incomplete"} { + t.Run(media+"/"+state, func(t *testing.T) { + f := newBrokerEvidenceFixture(t, func(request foundryResponseRequest) (string, []byte) { + response := map[string]any{"id": "provider-failure", "status": state, "agent_session_id": request.AgentSessionID, "output": []any{}} + if state == "failed" { + response["error"] = map[string]string{"code": "server_error", "message": "fixture-error-do-not-persist"} + } else { + response["incomplete_details"] = map[string]string{"reason": "max_output_tokens"} + } + if media == "text/event-stream" { + data, _ := json.Marshal(map[string]any{"type": "response." + state, "response": response}) + return media, []byte("data: " + string(data) + "\n\n") + } + data, _ := json.Marshal(response) + return media, data + }) + cfg := brokerTestConfig(t, f) + b, server := startBrokerTest(t, cfg) + c := brokerTestContext(cfg) + status, _, err := brokerTestHTTP(context.Background(), server.URL, brokerResponsesPath, c, brokerTestBody("")) + if err != nil || status == http.StatusOK { + t.Fatal("failed provider response was returned as success") + } + if actual := brokerInvocationState(b, c); actual != "accepted" && actual != "settled" { + t.Fatalf("coherent failure lost its acknowledgement: state=%s", actual) + } + proof := brokerTestControl(t, server.URL, brokerSettlePath, c) + if !proof.SettlementProven || proof.ActiveInvocations != 0 || proof.AmbiguousInvocations != 0 { + t.Fatal("acknowledged failure could not prove stop settlement") + } + proof = brokerTestControl(t, server.URL, brokerRetirePath, c) + if !proof.RetirementProven { + t.Fatal("acknowledged failure could not retire its owner") + } + creates, inferences, stops, deletes := f.counts() + if creates != 1 || inferences != 1 || stops == 0 || deletes != 1 { + t.Fatal("failure cleanup skipped containment or replayed work") + } + ledger, err := os.ReadFile(filepath.Join(cfg.stateDir, "state.json")) + if err != nil || bytes.Contains(ledger, []byte("fixture-error-do-not-persist")) { + t.Fatal("failure evidence persisted provider content") + } + }) + } + } +} + +func TestBrokerValidAcknowledgementSurvivesMalformedTail(t *testing.T) { + f := newBrokerEvidenceFixture(t, func(request foundryResponseRequest) (string, []byte) { + response := map[string]any{"id": "provider-accepted", "status": "in_progress", "agent_session_id": request.AgentSessionID, "output": []any{}} + data, _ := json.Marshal(map[string]any{"type": "response.created", "response": response}) + return "text/event-stream", []byte("data: " + string(data) + "\n\ndata: {\"type\":\"invalid-after-ack\"}\n\n") + }) + cfg := brokerTestConfig(t, f) + _, server := startBrokerTest(t, cfg) + c := brokerTestContext(cfg) + status, _, err := brokerTestHTTP(context.Background(), server.URL, brokerResponsesPath, c, brokerTestBody("")) + if err != nil || status == http.StatusOK { + t.Fatal("malformed tail exposed a successful response") + } + proof := brokerTestControl(t, server.URL, brokerSettlePath, c) + if !proof.SettlementProven || proof.AmbiguousInvocations != 0 { + t.Fatal("later malformed data erased valid acknowledgement") + } + _ = brokerTestControl(t, server.URL, brokerRetirePath, c) + creates, inferences, stops, deletes := f.counts() + if creates != 1 || inferences != 1 || stops == 0 || deletes != 1 { + t.Fatal("acknowledged malformed response lacked exact containment") + } +} + +type brokerEvidenceTokenProvider func(context.Context) (string, error) + +func (p brokerEvidenceTokenProvider) AccessToken(ctx context.Context) (string, error) { return p(ctx) } + +func TestBrokerUnsentRequestsRetainNoAmbiguousIntent(t *testing.T) { + for _, phase := range []string{"create", "inference"} { + for _, failure := range []string{"token-error", "malformed-token", "principal-drift"} { + t.Run(phase+"/"+failure, func(t *testing.T) { + f := newBrokerFixture(t, "success") + cfg := brokerTestConfig(t, f) + failAt := int64(3) // Two target validation GETs precede the create POST. + if phase == "inference" { + failAt = 5 // Creation and its exact-session GET precede inference. + } + var tokenCalls atomic.Int64 + provider := brokerEvidenceTokenProvider(func(context.Context) (string, error) { + if tokenCalls.Add(1) == failAt { + switch failure { + case "token-error": + return "", errors.New("fixture identity unavailable") + case "malformed-token": + return "invalid-fixture-identity", nil + case "principal-drift": + claims := []byte(`{"aud":"https://ai.azure.com","tid":"test-tenant","oid":"other-principal","appid":"test-client"}`) + return "fixture." + base64.RawURLEncoding.EncodeToString(claims) + ".fixture", nil + } + } + return brokerTestToken(), nil + }) + b, err := newLifecycleBroker(context.Background(), cfg, provider, nil) + if err != nil { + t.Fatal("fixture broker could not start") + } + server := httptest.NewServer(b) + t.Cleanup(func() { b.close(); server.Close() }) + c := brokerTestContext(cfg) + status, _, err := brokerTestHTTP(context.Background(), server.URL, brokerResponsesPath, c, brokerTestBody("")) + if err != nil || status == http.StatusOK || tokenCalls.Load() < failAt { + t.Fatal("pre-send identity failure was not exercised") + } + b.mu.Lock() + owner := b.ledger.Sessions[brokerJSONDigest(c.Owner)] + createState, remoteID := owner.CreateState, owner.RemoteID + b.mu.Unlock() + if phase == "create" && (createState != "none" || remoteID != "") { + t.Fatalf("unsent creation retained ambiguous ownership: state=%s", createState) + } + proof := brokerTestControl(t, server.URL, brokerSettlePath, c) + if !proof.SettlementProven || proof.CreatePending || proof.AmbiguousInvocations != 0 { + t.Fatal("definitively unsent request did not settle") + } + proof = brokerTestControl(t, server.URL, brokerRetirePath, c) + if !proof.RetirementProven { + t.Fatal("unsent request left an unretirable owner") + } + creates, inferences, stops, deletes := f.counts() + if inferences != 0 || (phase == "create" && (creates != 0 || stops != 0 || deletes != 0)) || + (phase == "inference" && (creates != 1 || stops != 0 || deletes != 1)) { + t.Fatal("pre-send failure replayed or dispatched remote work") + } + }) + } + } +} diff --git a/broker_folded_authority_test.go b/broker_folded_authority_test.go new file mode 100644 index 0000000..668e17d --- /dev/null +++ b/broker_folded_authority_test.go @@ -0,0 +1,122 @@ +package main + +import ( + "context" + "encoding/json" + "fmt" + "io" + "net/http" + "os" + "path/filepath" + "strings" + "sync/atomic" + "testing" + "time" +) + +func TestBrokerFoldedSessionEvidenceCannotSettleActiveOwner(t *testing.T) { + f := newBrokerFixture(t, "hold-known") + var cleanupEvidence atomic.Bool + transport := http.DefaultTransport.(*http.Transport).Clone() + t.Cleanup(transport.CloseIdleConnections) + client := &http.Client{Transport: brokerFixtureTransport(func(r *http.Request) (*http.Response, error) { + if r.Method == http.MethodPost && strings.HasSuffix(r.URL.Path, ":stop") { + // The original remote inference stays active. A 204 alone cannot + // prove idle, and the following GET is deliberately contradictory. + f.mu.Lock() + f.stops++ + f.mu.Unlock() + return &http.Response{StatusCode: http.StatusNoContent, Body: http.NoBody, Header: make(http.Header), Request: r}, nil + } + response, err := transport.RoundTrip(r) + if err != nil || !cleanupEvidence.Load() || r.Method != http.MethodGet || !strings.Contains(r.URL.Path, "/endpoint/sessions/") || response.StatusCode != http.StatusOK { + return response, err + } + data, err := io.ReadAll(response.Body) + _ = response.Body.Close() + if err != nil { + return nil, err + } + var session brokerRemoteSession + if json.Unmarshal(data, &session) != nil { + t.Error("fixture session evidence unreadable") + } + data = []byte(fmt.Sprintf(`{"agent_session_id":%q,"version_indicator":{"type":"version_ref","agent_version":"3"},"status":"active","Status":"idle"}`, session.ID)) + response.Body = io.NopCloser(strings.NewReader(string(data))) + response.ContentLength = int64(len(data)) + return response, nil + })} + cfg := brokerTestConfig(t, f) + b, server := startBrokerTestWithClient(t, cfg, client) + c := brokerTestContext(cfg) + key := brokerJSONDigest(c.Owner) + finished := make(chan struct{}) + go func() { + _, _, _ = brokerTestHTTP(context.Background(), server.URL, brokerResponsesPath, c, brokerTestBody("")) + close(finished) + }() + brokerAwait(t, func() bool { + b.mu.Lock() + defer b.mu.Unlock() + session := b.ledger.Sessions[key] + return session != nil && session.Prompts[c.promptKey()].Invocations[1].ResponseID != "" + }) + cleanupEvidence.Store(true) + control, body := brokerTestControlContext(brokerSettlePath, c) + status, _, err := brokerTestHTTP(context.Background(), server.URL, brokerSettlePath, control, body) + if err != nil || (status != http.StatusOK && status != http.StatusConflict) { + t.Fatalf("original settlement request failed: status=%d", status) + } + select { + case <-finished: + case <-time.After(3 * time.Second): + t.Fatal("original local invocation did not stop") + } + ctx, cancel := context.WithTimeout(context.Background(), time.Second) + defer cancel() + b.reconcile(ctx, key) + b.mu.Lock() + proof := b.controlResponseLocked(c, "") + remoteID := b.ledger.Sessions[key].RemoteID + b.mu.Unlock() + data, err := os.ReadFile(filepath.Join(cfg.stateDir, "state.json")) + var ledger brokerLedger + if err != nil || json.Unmarshal(data, &ledger) != nil { + t.Fatal("durable evidence unreadable") + } + f.mu.Lock() + active := f.sessions[remoteID] == "active" + f.mu.Unlock() + creates, inferences, stops, deletes := f.counts() + durable := ledger.Sessions[key].Prompts[c.promptKey()].Settled + if !active || creates != 1 || inferences != 1 || deletes != 0 { + t.Fatal("fixture did not preserve the one original active remote request") + } + if proof.SettlementProven || durable || proof.ProofDigest != "" { + t.Fatalf("malformed Session metadata proved cleanup: active=%v settled=%v durable=%v stops=%d", active, proof.SettlementProven, durable, stops) + } +} + +func TestBrokerFoldedSessionIdentityRejected(t *testing.T) { + for name, document := range map[string]string{ + "session identity": `{"agent_session_id":"wrong","Agent_Session_ID":"owned","version_indicator":{"type":"version_ref","agent_version":"3"},"status":"idle"}`, + "nested version": `{"agent_session_id":"owned","version_indicator":{"type":"version_ref","agent_version":"wrong","Agent_Version":"3"},"status":"idle"}`, + "nested type": `{"agent_session_id":"owned","version_indicator":{"type":"wrong","Type":"version_ref","agent_version":"3"},"status":"idle"}`, + "merged version object": `{"agent_session_id":"owned","version_indicator":{"type":"version_ref","agent_version":"wrong"},"Version_Indicator":{"agent_version":"3"},"status":"idle"}`, + "Unicode status": `{"agent_session_id":"owned","version_indicator":{"type":"version_ref","agent_version":"3"},"status":"active","ſtatus":"idle"}`, + "escaped status": `{"agent_session_id":"owned","version_indicator":{"type":"version_ref","agent_version":"3"},"status":"active","\u0053tatus":"idle"}`, + } { + t.Run(name, func(t *testing.T) { + f := newBrokerFixture(t, "success") + cfg := brokerTestConfig(t, f) + client := &http.Client{Transport: brokerFixtureTransport(func(r *http.Request) (*http.Response, error) { + return &http.Response{StatusCode: http.StatusOK, Header: http.Header{"Content-Type": []string{"application/json"}}, Body: io.NopCloser(strings.NewReader(document)), Request: r}, nil + })} + b, _ := startBrokerTestWithClient(t, cfg, client) + _, _, err := b.remoteSessionGet(context.Background(), "owned") + if err == nil { + t.Fatal("ambiguous remote identity or state accepted") + } + }) + } +} diff --git a/broker_guards_test.go b/broker_guards_test.go new file mode 100644 index 0000000..95a40c6 --- /dev/null +++ b/broker_guards_test.go @@ -0,0 +1,219 @@ +package main + +import ( + "bytes" + "context" + "encoding/base64" + "encoding/json" + "net/http" + "net/http/httptest" + "strings" + "testing" + "time" +) + +func brokerFunctionBody(previous, call string) []byte { + body, _ := json.Marshal(acpResponseRequest{Model: "fixture-model", foundryResponseRequest: foundryResponseRequest{ + Stream: true, Store: true, PreviousResponseID: previous, + Input: []foundryFunctionOutput{{Type: "function_call_output", CallID: call, Output: "fixture-tool-result"}}, + }}) + return body +} + +func TestBrokerOpaqueAliasesAndNoFunctionReplay(t *testing.T) { + f := newBrokerFixture(t, "functions") + cfg := brokerTestConfig(t, f) + _, server := startBrokerTest(t, cfg) + c := brokerTestContext(cfg) + c.InvocationSequence = 3 // Supervisor allocation order can contain gaps. + status, data, err := brokerTestHTTP(context.Background(), server.URL, brokerResponsesPath, c, brokerTestBody("")) + var first foundryResponse + if err != nil || status != 200 || json.Unmarshal(data, &first) != nil || len(first.Output) != 1 { + t.Fatalf("function proposal failed: %d", status) + } + call := first.Output[0].CallID + if !brokerAliasValid(first.ID, "fr_") || !brokerAliasValid(call, "fc_") || + !brokerAliasValid(first.Output[0].ID, "fi_") || bytes.Contains(data, []byte("provider-")) { + t.Fatal("native identities were not replaced by opaque aliases") + } + for _, which := range []string{"duplicate", "native_call", "cross_owner", "cross_prompt", "missing_output"} { + bad := c + bad.OperationID = "bad-" + which + bad.InvocationSequence = 4 + body := brokerFunctionBody(first.ID, call) + switch which { + case "duplicate": + bad = c + body = brokerTestBody("") + case "native_call": + body = brokerFunctionBody(first.ID, "provider-call-id") + case "cross_owner": + bad.Owner.RuntimeSessionUID = "another-runtime-session" + case "cross_prompt": + bad.PromptID = "another-prompt" + bad.TaskUID = "another-task" + case "missing_output": + body = brokerTestBody(first.ID) + } + status, _, err = brokerTestHTTP(context.Background(), server.URL, brokerResponsesPath, bad, body) + if err != nil || status != http.StatusConflict { + t.Fatalf("invalid function ownership %s was accepted: %d", which, status) + } + } + next := c + next.OperationID = "function-output" + next.InvocationSequence = 8 + status, data, err = brokerTestHTTP(context.Background(), server.URL, brokerResponsesPath, next, brokerFunctionBody(first.ID, call)) + var second foundryResponse + if err != nil || status != 200 || json.Unmarshal(data, &second) != nil { + t.Fatal("valid owned function output failed") + } + for _, previous := range []string{first.ID, second.ID} { + replay := next + replay.InvocationSequence = 9 + replay.OperationID = "replay-" + previous + status, _, err = brokerTestHTTP(context.Background(), server.URL, brokerResponsesPath, replay, brokerFunctionBody(previous, call)) + if err != nil || status != http.StatusConflict { + t.Fatal("consumed function output was replayed") + } + } + f.mu.Lock() + if len(f.requests) != 2 || f.requests[1].PreviousResponseID != "provider-response-1" { + t.Error("provider continuation did not use owned response identity") + } else { + items, ok := f.requests[1].Input.([]any) + if !ok || len(items) != 1 || items[0].(map[string]any)["call_id"] != "provider-call-id" { + t.Error("provider function result did not use owned call identity") + } + } + f.mu.Unlock() + _ = brokerTestControl(t, server.URL, brokerSettlePath, next) + _ = brokerTestControl(t, server.URL, brokerRetirePath, next) +} + +func TestBrokerConcurrentDuplicateAdmitsOnce(t *testing.T) { + f := newBrokerFixture(t, "hold-known") + cfg := brokerTestConfig(t, f) + b, server := startBrokerTest(t, cfg) + c := brokerTestContext(cfg) + first := brokerAsyncInference(context.Background(), server.URL, c, brokerTestBody("")) + brokerAwait(t, func() bool { return brokerInvocationState(b, c) == "accepted" }) + others := make([]<-chan brokerHTTPResult, 8) + for i := range others { + others[i] = brokerAsyncInference(context.Background(), server.URL, c, brokerTestBody("")) + } + for _, done := range others { + result := brokerWaitInference(t, done) + if result.err != nil || result.status != http.StatusConflict { + t.Fatal("concurrent duplicate was admitted") + } + } + f.unblock() + result := brokerWaitInference(t, first) + if result.err != nil || result.status != http.StatusOK { + t.Fatal("original concurrent inference did not complete") + } + creates, inferences, _, _ := f.counts() + if creates != 1 || inferences != 1 { + t.Fatal("duplicate caused a second provider operation") + } + _ = brokerTestControl(t, server.URL, brokerSettlePath, c) + _ = brokerTestControl(t, server.URL, brokerRetirePath, c) +} + +func TestBrokerRejectsUntrustedRoutesContextsAndInput(t *testing.T) { + f := newBrokerFixture(t, "success") + cfg := brokerTestConfig(t, f) + b, _ := startBrokerTest(t, cfg) + for _, name := range []string{"auth", "query", "duplicate_header", "malformed_context", "duplicate_context_member", + "config_digest", "empty_epoch", "wrong_digest", "expired", "model", "duplicate_body_member", "session", "background", + "conversation", "static_tools", "native_item_reference", "object_input", "nonstream", "nonstore"} { + t.Run(name, func(t *testing.T) { + c := brokerTestContext(cfg) + body := brokerTestBody("") + path := brokerResponsesPath + switch name { + case "query": + path += "?unexpected=1" + case "config_digest": + c.AgentConfigurationDigest = "sha256:" + strings.Repeat("3", 64) + case "empty_epoch": + c.Owner.ControllerEpoch = 0 + case "expired": + c.LeaseExpiresAt = time.Now().Add(-time.Second).UTC().Format(time.RFC3339Nano) + case "model": + body = bytes.Replace(body, []byte("fixture-model"), []byte("different-model"), 1) + case "duplicate_body_member": + body = append([]byte(`{"model":"fixture-model",`), body[1:]...) + case "session", "background", "conversation", "static_tools": + field := map[string]string{"session": `"agent_session_id":"injected",`, "background": `"background":true,`, + "conversation": `"conversation":"injected",`, "static_tools": `"tools":[],`}[name] + body = append([]byte("{"+field), body[1:]...) + case "native_item_reference", "object_input": + var fields map[string]any + _ = json.Unmarshal(body, &fields) + fields["input"] = map[string]any{"type": "item_reference", "id": "injected-provider-id"} + if name == "native_item_reference" { + fields["input"] = []any{fields["input"]} + } + body, _ = json.Marshal(fields) + case "nonstream": + body = bytes.Replace(body, []byte(`"stream":true`), []byte(`"stream":false`), 1) + case "nonstore": + body = bytes.Replace(body, []byte(`"store":true`), []byte(`"store":false`), 1) + } + c.BodySHA256 = brokerSHA(body) + if name == "wrong_digest" { + c.BodySHA256 = "sha256:" + strings.Repeat("3", 64) + } + raw, _ := json.Marshal(c) + if name == "duplicate_context_member" { + raw = append([]byte(`{"protocol":"orka.foundry.broker.v1",`), raw[1:]...) + } + request := httptest.NewRequest(http.MethodPost, path, bytes.NewReader(body)) + request.Header.Set("Authorization", "Bearer "+brokerFixtureBearer) + request.Header.Set(brokerContextHeader, base64.RawURLEncoding.EncodeToString(raw)) + if name == "auth" { + request.Header.Del("Authorization") + } + if name == "malformed_context" { + request.Header.Set(brokerContextHeader, "invalid==") + } + if name == "duplicate_header" { + request.Header.Add(brokerContextHeader, request.Header.Get(brokerContextHeader)) + } + response := httptest.NewRecorder() + b.ServeHTTP(response, request) + if response.Code != http.StatusBadRequest && response.Code != http.StatusGone && response.Code != http.StatusUnauthorized { + t.Fatalf("invalid request was accepted: %d", response.Code) + } + }) + } + creates, inferences, stops, deletes := f.counts() + if creates+inferences+stops+deletes != 0 { + t.Fatal("an invalid request reached provider mutations") + } +} + +func TestBrokerControlProofBindsExactHeaderBytes(t *testing.T) { + f := newBrokerFixture(t, "success") + cfg := brokerTestConfig(t, f) + _, server := startBrokerTest(t, cfg) + c, body := brokerTestControlContext(brokerSettlePath, brokerTestContext(cfg)) + c.BodySHA256 = brokerSHA(body) + status, data, err := brokerTestHTTP(context.Background(), server.URL, brokerSettlePath, c, body) + var proof brokerControlResponse + if err != nil || (status != 200 && status != 409) || json.Unmarshal(data, &proof) != nil { + t.Fatal("control reply unavailable") + } + if proof.Protocol != brokerProtocol || proof.OwnerDigest != brokerJSONDigest(c.Owner) || + proof.OperationID != c.OperationID || proof.ContextSHA256 != brokerJSONDigest(c) { + t.Fatal("proof did not bind the exact trusted owner and context") + } + // An operation ID cannot be reassigned to a different owner-context body. + c.LeaseGeneration++ + status, _, err = brokerTestHTTP(context.Background(), server.URL, brokerSettlePath, c, body) + if err != nil || status != http.StatusConflict { + t.Fatal("control idempotency key accepted another context") + } +} diff --git a/broker_lease_admission_test.go b/broker_lease_admission_test.go new file mode 100644 index 0000000..559e7d7 --- /dev/null +++ b/broker_lease_admission_test.go @@ -0,0 +1,60 @@ +package main + +import ( + "bytes" + "context" + "net/http" + "os" + "path/filepath" + "testing" + "time" +) + +func TestBrokerNewInferenceRequiresCurrentLease(t *testing.T) { + for _, kind := range []string{"superseded-generation", "superseded-expiry", "current"} { + t.Run(kind, func(t *testing.T) { + f := newBrokerFixture(t, "success") + cfg := brokerTestConfig(t, f) + _, server := startBrokerTest(t, cfg) + c := brokerTestContext(cfg) + first, body := brokerTestControlContext(brokerRenewPath, c) + first.OperationID = "original-lease" + status, _, err := brokerTestHTTP(context.Background(), server.URL, brokerRenewPath, first, body) + if err != nil || status != http.StatusOK { + t.Fatal("could not establish original lease") + } + renewed := first + renewed.OperationID = "renewed-lease" + renewed.LeaseGeneration++ + renewed.LeaseExpiresAt = time.Now().Add(20 * time.Second).UTC().Format(time.RFC3339Nano) + status, _, err = brokerTestHTTP(context.Background(), server.URL, brokerRenewPath, renewed, body) + if err != nil || status != http.StatusOK { + t.Fatal("exact renewal did not acknowledge") + } + if kind != "superseded-generation" { + c.LeaseGeneration = renewed.LeaseGeneration + } + if kind != "superseded-expiry" { + c.LeaseExpiresAt = renewed.LeaseExpiresAt + } + c.OperationID = "fresh-inference" + before, err := os.ReadFile(filepath.Join(cfg.stateDir, "state.json")) + if err != nil { + t.Fatal("could not read original ownership") + } + status, _, err = brokerTestHTTP(context.Background(), server.URL, brokerResponsesPath, c, brokerTestBody("")) + creates, inferences, _, _ := f.counts() + if err != nil { + t.Fatal("broker response unreadable") + } + if kind != "current" { + after, readErr := os.ReadFile(filepath.Join(cfg.stateDir, "state.json")) + if status != http.StatusConflict || creates != 0 || inferences != 0 || readErr != nil || !bytes.Equal(before, after) { + t.Fatalf("superseded lease admitted new work: status=%d creates=%d inferences=%d", status, creates, inferences) + } + } else if status != http.StatusOK || creates != 1 || inferences != 1 { + t.Fatalf("current lease rejected: status=%d creates=%d inferences=%d", status, creates, inferences) + } + }) + } +} diff --git a/broker_legacy_recovery_test.go b/broker_legacy_recovery_test.go new file mode 100644 index 0000000..e224be8 --- /dev/null +++ b/broker_legacy_recovery_test.go @@ -0,0 +1,63 @@ +package main + +import ( + "bytes" + "context" + "encoding/json" + "net/http" + "os" + "path/filepath" + "testing" + "time" +) + +func TestBrokerLegacyIntentAtByteCapRemainsContainable(t *testing.T) { + f := newBrokerFixture(t, "hold-unknown") + cfg := brokerTestConfig(t, f) + b, server := startBrokerTest(t, cfg) + c := brokerTestContext(cfg) + c.LeaseExpiresAt = time.Now().Add(4 * time.Minute).UTC().Format(time.RFC3339Nano) + done := brokerAsyncInference(context.Background(), server.URL, c, brokerTestBody("")) + select { + case <-f.started: + case <-time.After(3 * time.Second): + t.Fatal("original request did not reach provider transport") + } + // The legacy schema admitted this valid retained history before ordinary + // admission reserved cleanup bytes. Preserve the actual in-flight owner. + brokerCapacityFillHistory(t, b, c, brokerMaxLedgerBytes, false) + path := filepath.Join(cfg.stateDir, "state.json") + before, err := os.ReadFile(path) + var original brokerLedger + if err != nil || json.Unmarshal(before, &original) != nil || !brokerLedgerValid(&original, cfg.configDigest) || len(before) != brokerMaxLedgerBytes { + t.Fatal("legacy fixture does not reach the valid byte boundary") + } + b.close() + server.Close() + result := brokerWaitInference(t, done) + if result.err == nil && result.status == http.StatusOK { + t.Fatal("interrupted inference returned usable output") + } + retained, err := os.ReadFile(path) + if err != nil || !bytes.Equal(before, retained) { + t.Fatal("unpersistable uncertainty overwrote the legacy owner") + } + recovered, restarted := startBrokerTest(t, cfg) + proof := brokerTestControl(t, restarted.URL, brokerStatusPath, c) + if proof.State != "blocked" || proof.AmbiguousInvocations != 1 || proof.ActiveInvocations != 0 || + proof.CreatePending || !proof.RemoteSessionCreated || proof.SettlementProven || proof.RetirementProven || proof.ProofDigest != "" { + t.Fatal("abandoned intent was reported as active or proven cleanup") + } + brokerAwait(t, func() bool { _, _, stops, _ := f.counts(); return stops > 0 }) + recovered.mu.Lock() + session := recovered.ledger.Sessions[brokerJSONDigest(c.Owner)] + prompt := session.Prompts[c.promptKey()] + valid := brokerLedgerValid(recovered.ledger, cfg.configDigest) && prompt.Closing && !prompt.Settled && + !session.Retired && prompt.Invocations[c.InvocationSequence].State == "intent" + recovered.mu.Unlock() + after, err := os.ReadFile(path) + creates, inferences, _, deletes := f.counts() + if !valid || err != nil || len(after) > len(before) || creates != 1 || inferences != 1 || deletes != 0 { + t.Fatal("legacy containment grew the ledger, lost ownership, replayed work, or claimed deletion") + } +} diff --git a/broker_main.go b/broker_main.go new file mode 100644 index 0000000..aa5697b --- /dev/null +++ b/broker_main.go @@ -0,0 +1,138 @@ +package main + +import ( + "context" + "errors" + "flag" + "io" + "net" + "net/http" + "os" + "os/signal" + "strconv" + "strings" + "syscall" + "time" +) + +type brokerConfiguration struct { + agent acpAgentConfiguration + configDigest string + addr string + stateDir string + bearer string + operationTimeout time.Duration +} + +func maybeServeBroker(args []string) (bool, error) { + selected := false + for index, arg := range args { + if arg == "--protocol=broker" || (arg == "--protocol" && index+1 < len(args) && args[index+1] == "broker") { + selected = true + } + } + if !selected { + return false, nil + } + flags := flag.NewFlagSet("foundry-broker", flag.ContinueOnError) + flags.SetOutput(io.Discard) + protocol := flags.String("protocol", "", "") + path := flags.String("config", acpConfigPath, "") + healthCheck := flags.Bool("health-check", false, "") + if flags.Parse(args) != nil || flags.NArg() != 0 || *protocol != "broker" { + return true, errBrokerInvalid + } + if *healthCheck { + return true, checkBrokerHealth(firstNonBlank(os.Getenv("ORKA_FOUNDRY_BROKER_ADDR"), "127.0.0.1:8091")) + } + cfg, err := loadBrokerConfiguration(*path, os.Getenv) + if err != nil { + return true, err + } + provider, err := newAzureFoundryTokenProvider() + if err != nil { + return true, errBrokerRemote + } + ctx, stop := signal.NotifyContext(context.Background(), os.Interrupt, syscall.SIGTERM) + defer stop() + broker, err := newLifecycleBroker(ctx, cfg, provider, nil) + if err != nil { + return true, err + } + defer broker.close() + 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) + go func() { finished <- server.ListenAndServe() }() + select { + case err := <-finished: + if errors.Is(err, http.ErrServerClosed) { + return true, nil + } + return true, errBrokerRemote + case <-ctx.Done(): + shutdown, cancel := context.WithTimeout(context.Background(), 5*time.Second) + defer cancel() + _ = server.Shutdown(shutdown) + return true, nil + } +} + +func loadBrokerConfiguration(path string, getenv func(string) string) (brokerConfiguration, error) { + file, err := os.Open(path) + if err != nil { + return brokerConfiguration{}, errBrokerInvalid + } + defer file.Close() //nolint:errcheck + data, err := io.ReadAll(io.LimitReader(file, acpMaxConfigBytes+1)) + if err != nil || len(data) > acpMaxConfigBytes { + return brokerConfiguration{}, errBrokerInvalid + } + digest := getenv(acpConfigDigestEnv) + agent, err := decodeACPAgentConfiguration(data, digest, getenv(acpModelEnv)) + if err != nil { + return brokerConfiguration{}, errBrokerInvalid + } + cfg := brokerConfiguration{agent: agent, configDigest: digest, + addr: 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"), + operationTimeout: 45 * time.Second} + if !brokerAddressValid(cfg.addr) || cfg.stateDir == "" || + !acpSafeString(cfg.bearer, 16<<10) || len(cfg.bearer) < 32 || strings.ContainsAny(cfg.bearer, " \t") || + (getenv(envIsolationMode) != "" && getenv(envIsolationMode) != "entra") { + return brokerConfiguration{}, errBrokerInvalid + } + return cfg, nil +} + +func brokerAddressValid(address string) bool { + host, port, err := net.SplitHostPort(address) + ip := net.ParseIP(host) + value, portErr := strconv.Atoi(port) + return err == nil && ip != nil && ip.IsLoopback() && portErr == nil && value > 0 && value <= 65535 +} + +// The distroless image can run its own exec readiness probe without credentials, +// a configuration file, or a second writer of the durable ledger. +func checkBrokerHealth(address string) error { + if !brokerAddressValid(address) { + return errBrokerInvalid + } + ctx, cancel := context.WithTimeout(context.Background(), 2*time.Second) + defer cancel() + request, err := http.NewRequestWithContext(ctx, http.MethodGet, "http://"+address+"/healthz", nil) + if err != nil { + return errBrokerInvalid + } + client := newBrokerHTTPClient() + defer client.CloseIdleConnections() + response, err := client.Do(request) + if err != nil { + return errBrokerRemote + } + defer response.Body.Close() //nolint:errcheck + if response.StatusCode != http.StatusOK { + return errBrokerRemote + } + return nil +} diff --git a/broker_preflight_recovery_test.go b/broker_preflight_recovery_test.go new file mode 100644 index 0000000..bf0ddf1 --- /dev/null +++ b/broker_preflight_recovery_test.go @@ -0,0 +1,156 @@ +package main + +import ( + "context" + "encoding/base64" + "encoding/json" + "errors" + "net/http" + "net/http/httptest" + "os" + "path/filepath" + "strings" + "sync/atomic" + "testing" +) + +func TestBrokerPreparedRequestsKeepDurableIntentBeforeTransport(t *testing.T) { + f := newBrokerFixture(t, "success") + cfg := brokerTestConfig(t, f) + c := brokerTestContext(cfg) + var posts atomic.Int64 + client := &http.Client{Transport: brokerFixtureTransport(func(request *http.Request) (*http.Response, error) { + if request.Method == http.MethodPost && (strings.HasSuffix(request.URL.Path, "/endpoint/sessions") || + strings.HasSuffix(request.URL.Path, "/endpoint/protocols/openai/responses")) { + raw, err := os.ReadFile(filepath.Join(cfg.stateDir, "state.json")) + var ledger brokerLedger + if err != nil || json.Unmarshal(raw, &ledger) != nil { + t.Error("outbound submission has no readable durable ownership") + return nil, errBrokerStorage + } + owner := ledger.Sessions[brokerJSONDigest(c.Owner)] + if owner == nil || owner.RemoteID == "" { + t.Error("outbound submission lost its exact owner") + return nil, errBrokerStorage + } + if strings.HasSuffix(request.URL.Path, "/endpoint/sessions") { + if owner.CreateState != "intent" { + t.Error("creation reached transport before durable intent") + } + } else if owner.Prompts[c.promptKey()].Invocations[c.InvocationSequence].State != "intent" { + t.Error("inference reached transport before durable intent") + } + if request.GetBody != nil { + t.Error("prepared request permits implicit replay") + } + posts.Add(1) + } + return http.DefaultTransport.RoundTrip(request) + })} + _, server := startBrokerTestWithClient(t, cfg, client) + status, _, err := brokerTestHTTP(context.Background(), server.URL, brokerResponsesPath, c, brokerTestBody("")) + if err != nil || status != http.StatusOK || posts.Load() != 2 { + t.Fatal("prepared request did not complete exactly one creation and inference") + } + _ = brokerTestControl(t, server.URL, brokerSettlePath, c) + _ = brokerTestControl(t, server.URL, brokerRetirePath, c) +} + +func TestBrokerPreflightCrashDoesNotStrandUnsentOwnership(t *testing.T) { + for _, phase := range []string{"create", "inference"} { + for _, failure := range []string{"token-error", "malformed-token", "principal-drift", "cancelled-token", "cancelled-after-token"} { + t.Run(phase+"/"+failure, func(t *testing.T) { + f := newBrokerFixture(t, "success") + cfg := brokerTestConfig(t, f) + c := brokerTestContext(cfg) + failAt := int64(3) + if phase == "inference" { + failAt = 5 + } + var calls atomic.Int64 + var b *lifecycleBroker + crashLedger := make(chan []byte, 1) + provider := brokerEvidenceTokenProvider(func(context.Context) (string, error) { + if calls.Add(1) != failAt { + return brokerTestToken(), nil + } + // Capture exactly what a restart can recover if the process + // dies during authentication, before any failure rollback. + raw, err := os.ReadFile(filepath.Join(cfg.stateDir, "state.json")) + if err != nil { + t.Error("could not capture preflight crash ledger") + return "", errBrokerRemote + } + crashLedger <- raw + switch failure { + case "token-error": + return "", errors.New("fixture authentication unavailable") + case "malformed-token": + return "invalid-fixture-identity", nil + case "principal-drift": + claims := []byte(`{"aud":"https://ai.azure.com","tid":"test-tenant","oid":"other-principal","appid":"test-client"}`) + return "fixture." + base64.RawURLEncoding.EncodeToString(claims) + ".fixture", nil + default: + b.mu.Lock() + cancel := b.active[brokerJSONDigest(c.Owner)].cancel + b.mu.Unlock() + cancel() + if failure == "cancelled-token" { + return "", context.Canceled + } + return brokerTestToken(), nil + } + }) + var err error + b, err = newLifecycleBroker(context.Background(), cfg, provider, nil) + if err != nil { + t.Fatal("could not start preflight crash fixture") + } + server := httptest.NewServer(b) + t.Cleanup(func() { b.close(); server.Close() }) + status, _, err := brokerTestHTTP(context.Background(), server.URL, brokerResponsesPath, c, brokerTestBody("")) + if err != nil || status == http.StatusOK { + t.Fatal("definitely-unsent preflight unexpectedly succeeded") + } + var raw []byte + select { + case raw = <-crashLedger: + default: + t.Fatal("request did not reach the selected authentication boundary") + } + b.close() + server.Close() + var ledger brokerLedger + if json.Unmarshal(raw, &ledger) != nil || !brokerLedgerValid(&ledger, cfg.configDigest) { + t.Fatal("preflight snapshot is not valid durable ownership") + } + owner := ledger.Sessions[brokerJSONDigest(c.Owner)] + if owner == nil || owner.Prompts[c.promptKey()] == nil { + t.Fatal("preflight snapshot lost reserved ownership") + } + if phase == "create" && (owner.CreateState != "none" || owner.RemoteID != "") { + t.Fatal("crash during create authentication leaves possibly-sent creation") + } + if owner.Prompts[c.promptKey()].Invocations[c.InvocationSequence].State != "reserved" { + t.Fatal("crash during authentication leaves possibly-sent inference") + } + // Discard later graceful-cleanup writes to model this exact crash. + if os.WriteFile(filepath.Join(cfg.stateDir, "state.json"), raw, 0o600) != nil { + t.Fatal("could not restore crash-boundary fixture") + } + _, restarted := startBrokerTest(t, cfg) + proof := brokerTestControl(t, restarted.URL, brokerSettlePath, c) + if !proof.SettlementProven || proof.CreatePending || proof.AmbiguousInvocations != 0 { + t.Fatal("restart stranded definitely-unsent ownership") + } + proof = brokerTestControl(t, restarted.URL, brokerRetirePath, c) + creates, inferences, stops, deletes := f.counts() + if !proof.RetirementProven || inferences != 0 || stops != 0 || + (phase == "create" && (creates != 0 || deletes != 0)) || + (phase == "inference" && (creates != 1 || deletes != 1)) { + t.Fatal("preflight crash recovery replayed work or lost exact retirement") + } + }) + } + } +} diff --git a/broker_protocol.go b/broker_protocol.go new file mode 100644 index 0000000..abe045a --- /dev/null +++ b/broker_protocol.go @@ -0,0 +1,160 @@ +package main + +import ( + "crypto/sha256" + "encoding/base64" + "encoding/hex" + "encoding/json" + "errors" + "net/http" + "strings" + "time" +) + +const ( + brokerProtocol = "orka.foundry.broker.v1" + brokerContextHeader = "X-Orka-Foundry-Context" + brokerMaxContext = 16 << 10 + brokerResponsesPath = "/v1/responses" + brokerRenewPath = "/internal/v1/renew" + brokerSettlePath = "/internal/v1/settle" + brokerRetirePath = "/internal/v1/retire" + brokerStatusPath = "/internal/v1/status" +) + +var ( + errBrokerInvalid = errors.New("invalid Foundry broker request") + errBrokerConflict = errors.New("Foundry broker operation conflicts with durable ownership") + errBrokerClosed = errors.New("Foundry broker authority is closed") + errBrokerPending = errors.New("Foundry broker remote settlement is pending") + errBrokerStorage = errors.New("Foundry broker durable state is unavailable") + errBrokerRemote = errors.New("Foundry broker remote operation failed") + errBrokerAmbiguous = errors.New("Foundry broker remote acceptance is unknown") +) + +// Keep this field order and tags identical to Orka's v2.Fence. Both sides hash +// json.Marshal(Fence), rather than relying on untrusted JSON member order. +type brokerOwner struct { + RuntimeInstanceID string `json:"runtimeInstanceID"` + SupervisorBootID string `json:"supervisorBootID"` + ControllerEpoch uint64 `json:"controllerEpoch"` + RuntimePoolUID string `json:"runtimePoolUID"` + RuntimePoolGeneration uint64 `json:"runtimePoolGeneration"` + RuntimeSessionUID string `json:"runtimeSessionUID,omitempty"` + RuntimeSessionGeneration uint64 `json:"runtimeSessionGeneration,omitempty"` + RuntimeProfileDigest string `json:"runtimeProfileDigest"` + ProfileDigestSchemaVersion uint32 `json:"profileDigestSchemaVersion"` +} + +type brokerContext struct { + Protocol string `json:"protocol"` + Owner brokerOwner `json:"owner"` + AgentConfigurationDigest string `json:"agentConfigurationDigest"` + TaskUID string `json:"taskUID,omitempty"` + TaskAttempt uint32 `json:"taskAttempt,omitempty"` + PromptID string `json:"promptID,omitempty"` + PromptRequestDigest string `json:"promptRequestDigest,omitempty"` + LeaseGeneration uint64 `json:"leaseGeneration,omitempty"` + LeaseExpiresAt string `json:"leaseExpiresAt,omitempty"` + OperationID string `json:"operationID"` + InvocationSequence uint64 `json:"invocationSequence,omitempty"` + BodySHA256 string `json:"bodySHA256"` +} + +type brokerControlResponse struct { + Protocol string `json:"protocol"` + OwnerDigest string `json:"ownerDigest"` + OperationID string `json:"operationID"` + ContextSHA256 string `json:"contextSHA256"` + State string `json:"state"` + SettlementProven bool `json:"settlementProven"` + RetirementProven bool `json:"retirementProven"` + ActiveInvocations uint32 `json:"activeInvocations"` + AmbiguousInvocations uint32 `json:"ambiguousInvocations"` + CreatePending bool `json:"createPending"` + RemoteSessionCreated bool `json:"remoteSessionCreated"` + LeaseGeneration uint64 `json:"leaseGeneration"` + LeaseExpiresAt string `json:"leaseExpiresAt"` + ProofDigest string `json:"proofDigest"` +} + +func brokerSHA(data []byte) string { + digest := sha256.Sum256(data) + return "sha256:" + hex.EncodeToString(digest[:]) +} + +func brokerJSONDigest(value any) string { + data, _ := json.Marshal(value) // Only concrete, JSON-safe broker structs are used. + return brokerSHA(data) +} + +func brokerDigestValid(value string) bool { + if len(value) != len("sha256:")+64 || !strings.HasPrefix(value, "sha256:") { + return false + } + for _, char := range value[len("sha256:"):] { + if !(char >= '0' && char <= '9') && !(char >= 'a' && char <= 'f') { + return false + } + } + return true +} + +func (o brokerOwner) valid() bool { + return acpSafeString(o.RuntimeInstanceID, 512) && acpSafeString(o.SupervisorBootID, 512) && + o.ControllerEpoch > 0 && acpSafeString(o.RuntimePoolUID, 512) && o.RuntimePoolGeneration > 0 && + acpSafeString(o.RuntimeSessionUID, 512) && o.RuntimeSessionGeneration > 0 && + brokerDigestValid(o.RuntimeProfileDigest) && o.ProfileDigestSchemaVersion == 1 +} + +func (c brokerContext) promptKey() string { + return brokerJSONDigest(struct { + TaskUID string `json:"taskUID"` + TaskAttempt uint32 `json:"taskAttempt"` + PromptID string `json:"promptID"` + PromptRequestDigest string `json:"promptRequestDigest"` + }{c.TaskUID, c.TaskAttempt, c.PromptID, c.PromptRequestDigest}) +} + +func brokerParseContext(r *http.Request, body []byte, configDigest string, now time.Time) (brokerContext, string, error) { + values := r.Header.Values(brokerContextHeader) + if len(values) != 1 || len(values[0]) == 0 || len(values[0]) > brokerMaxContext { + return brokerContext{}, "", errBrokerInvalid + } + raw, err := base64.RawURLEncoding.Strict().DecodeString(values[0]) + if err != nil || base64.RawURLEncoding.EncodeToString(raw) != values[0] { + return brokerContext{}, "", errBrokerInvalid + } + var c brokerContext + if acpDecode(raw, &c, true) != nil || c.Protocol != brokerProtocol || !c.Owner.valid() || + c.AgentConfigurationDigest != configDigest || !brokerDigestValid(configDigest) || + !acpSafeString(c.OperationID, 512) || c.BodySHA256 != brokerSHA(body) { + return brokerContext{}, "", errBrokerInvalid + } + needsPrompt := r.URL.Path != brokerRetirePath && r.URL.Path != brokerStatusPath + if needsPrompt { + expiry, err := time.Parse(time.RFC3339Nano, c.LeaseExpiresAt) + if !acpSafeString(c.TaskUID, 512) || c.TaskAttempt == 0 || !acpSafeString(c.PromptID, 512) || + !brokerDigestValid(c.PromptRequestDigest) || c.LeaseGeneration == 0 || err != nil { + return brokerContext{}, "", errBrokerInvalid + } + if (r.URL.Path == brokerResponsesPath || r.URL.Path == brokerRenewPath) && + (!expiry.After(now) || expiry.After(now.Add(5*time.Minute))) { + return brokerContext{}, "", errBrokerClosed + } + } else if c.TaskUID != "" || c.TaskAttempt != 0 || c.PromptID != "" || c.PromptRequestDigest != "" || + c.LeaseGeneration != 0 || c.LeaseExpiresAt != "" { + return brokerContext{}, "", errBrokerInvalid + } + if (r.URL.Path == brokerResponsesPath) != (c.InvocationSequence > 0) { + return brokerContext{}, "", errBrokerInvalid + } + return c, brokerSHA(raw), nil +} + +func brokerOperationDigest(path string, c brokerContext) string { + return brokerJSONDigest(struct { + Path string `json:"path"` + Context brokerContext `json:"context"` + }{path, c}) +} diff --git a/broker_recovery_test.go b/broker_recovery_test.go new file mode 100644 index 0000000..bfe4341 --- /dev/null +++ b/broker_recovery_test.go @@ -0,0 +1,366 @@ +package main + +import ( + "context" + "encoding/json" + "net/http" + "os" + "path/filepath" + "testing" + "time" +) + +type brokerHTTPResult struct { + status int + data []byte + err error +} + +func brokerAsyncInference(ctx context.Context, base string, c brokerContext, body []byte) <-chan brokerHTTPResult { + done := make(chan brokerHTTPResult, 1) + go func() { + status, data, err := brokerTestHTTP(ctx, base, brokerResponsesPath, c, body) + done <- brokerHTTPResult{status, data, err} + }() + return done +} + +func brokerWaitInference(t *testing.T, done <-chan brokerHTTPResult) brokerHTTPResult { + t.Helper() + select { + case result := <-done: + return result + case <-time.After(4 * time.Second): + t.Fatal("inference did not settle within bound") + return brokerHTTPResult{} + } +} + +func brokerInvocationState(b *lifecycleBroker, c brokerContext) string { + b.mu.Lock() + defer b.mu.Unlock() + if session := b.ledger.Sessions[brokerJSONDigest(c.Owner)]; session != nil { + if prompt := session.Prompts[c.promptKey()]; prompt != nil { + if invocation := prompt.Invocations[c.InvocationSequence]; invocation != nil { + return invocation.State + } + } + } + return "" +} + +func brokerPendingControl(t *testing.T, base, path string, c brokerContext, pendingCreate bool, ambiguous uint32) { + t.Helper() + cc, body := brokerTestControlContext(path, c) + status, data, err := brokerTestHTTP(context.Background(), base, path, cc, body) + var proof brokerControlResponse + if err != nil || status != http.StatusConflict || json.Unmarshal(data, &proof) != nil || + proof.CreatePending != pendingCreate || proof.AmbiguousInvocations != ambiguous || + proof.SettlementProven || proof.RetirementProven || proof.ProofDigest != "" { + t.Fatalf("pending operation claimed proof or lost uncertainty: status=%d", status) + } +} + +func TestBrokerAcknowledgedDisconnectExpiryAndTruncation(t *testing.T) { + for _, mode := range []string{"disconnect", "expiry", "truncated"} { + t.Run(mode, func(t *testing.T) { + fixtureMode := "hold-known" + if mode == "truncated" { + fixtureMode = mode + } + f := newBrokerFixture(t, fixtureMode) + cfg := brokerTestConfig(t, f) + b, server := startBrokerTest(t, cfg) + c := brokerTestContext(cfg) + if mode == "expiry" { + c.LeaseExpiresAt = time.Now().Add(700 * time.Millisecond).UTC().Format(time.RFC3339Nano) + } + ctx, cancel := context.WithCancel(context.Background()) + defer cancel() + done := brokerAsyncInference(ctx, server.URL, c, brokerTestBody("")) + brokerAwait(t, func() bool { + state := brokerInvocationState(b, c) + return state == "accepted" || state == "settled" + }) + if mode == "disconnect" { + cancel() + } + result := brokerWaitInference(t, done) + if result.err == nil && result.status == http.StatusOK { + t.Fatal("interrupted inference exposed a terminal result") + } + proof := brokerTestControl(t, server.URL, brokerSettlePath, c) + if !proof.SettlementProven || proof.ActiveInvocations != 0 || proof.AmbiguousInvocations != 0 { + t.Fatal("acknowledged interrupted response did not settle") + } + creates, inferences, stops, deletes := f.counts() + if creates != 1 || inferences != 1 || stops < 1 || deletes != 0 { + t.Fatal("disconnect/expiry cleanup skipped stop or replayed inference") + } + _ = brokerTestControl(t, server.URL, brokerRetirePath, c) + }) + } +} + +func TestBrokerDelayedCreateRetainsOwnershipThrough404(t *testing.T) { + f := newBrokerFixture(t, "late-create") + cfg := brokerTestConfig(t, f) + b, server := startBrokerTest(t, cfg) + c := brokerTestContext(cfg) + ctx, cancel := context.WithCancel(context.Background()) + defer cancel() + done := brokerAsyncInference(ctx, server.URL, c, brokerTestBody("")) + select { + case <-f.started: + case <-time.After(3 * time.Second): + t.Fatal("create did not start") + } + cancel() + _ = brokerWaitInference(t, done) + brokerAwait(t, func() bool { return brokerInvocationState(b, c) == "rejected" }) + for range 3 { + brokerPendingControl(t, server.URL, brokerRetirePath, c, true, 0) + } + creates, inferences, _, deletes := f.counts() + if creates != 1 || inferences != 0 || deletes != 0 { + t.Fatal("unknown creation retried or deleted without acceptance proof") + } + // Even a broker restart must retain the original ID and its pending intent. + b.close() + server.Close() + _, server = startBrokerTest(t, cfg) + brokerPendingControl(t, server.URL, brokerRetirePath, c, true, 0) + f.unblock() + brokerAwait(t, func() bool { + f.mu.Lock() + defer f.mu.Unlock() + return len(f.sessions) == 1 + }) + // The object appearing after a lost acknowledgement does not establish + // completion of the original CREATE, including after a broker restart. + brokerPendingControl(t, server.URL, brokerRetirePath, c, true, 0) + creates, inferences, _, deletes = f.counts() + if creates != 1 || inferences != 0 || deletes != 0 { + t.Fatal("same-intent recovery replayed work or deleted an unacknowledged owner") + } +} + +func TestBrokerCancellationWaitsForCreateAcknowledgement(t *testing.T) { + f := newBrokerFixture(t, "hold-create") + cfg := brokerTestConfig(t, f) + providerClient := newBrokerHTTPClient() + transport := providerClient.Transport + creationContexts := make(chan context.Context, 1) + providerClient.Transport = brokerFixtureTransport(func(r *http.Request) (*http.Response, error) { + if r.Method == http.MethodPost && r.URL.Path == "/api/projects/fixture/agents/fixture/endpoint/sessions" { + select { + case creationContexts <- r.Context(): + default: + t.Error("session creation was replayed") + } + } + return transport.RoundTrip(r) + }) + _, server := startBrokerTestWithClient(t, cfg, providerClient) + c := brokerTestContext(cfg) + ctx, cancel := context.WithCancel(context.Background()) + defer cancel() + done := brokerAsyncInference(ctx, server.URL, c, brokerTestBody("")) + select { + case <-f.started: + case <-time.After(3 * time.Second): + t.Fatal("creation acknowledgement was not held") + } + createCtx := <-creationContexts + cancel() + if result := brokerWaitInference(t, done); result.err == nil { + t.Fatal("cancelled caller received a response") + } + // This public close request synchronously cancels the prompt context while + // the fixture still holds the original creation acknowledgement. + brokerPendingControl(t, server.URL, brokerSettlePath, c, true, 0) + if createCtx.Err() != nil { + t.Fatal("prompt cancellation aborted durable session creation") + } + if deadline, ok := createCtx.Deadline(); !ok || time.Until(deadline) > cfg.operationTimeout { + t.Fatal("session creation lacks the broker's operation bound") + } + f.unblock() + proof := brokerTestControl(t, server.URL, brokerSettlePath, c) + if !proof.SettlementProven || !proof.RemoteSessionCreated || proof.CreatePending || + proof.ActiveInvocations != 0 || proof.AmbiguousInvocations != 0 { + t.Fatal("acknowledged creation did not settle the cancelled prompt") + } + data, err := os.ReadFile(filepath.Join(cfg.stateDir, "state.json")) + var ledger brokerLedger + if err != nil || json.Unmarshal(data, &ledger) != nil || !brokerLedgerValid(&ledger, cfg.configDigest) { + t.Fatal("settlement did not preserve a valid durable ledger") + } + owned := ledger.Sessions[brokerJSONDigest(c.Owner)] + if owned == nil || owned.CreateState != "known" || owned.RemoteID == "" || !owned.Prompts[c.promptKey()].Settled { + t.Fatal("positive creation acknowledgement was not durably retained") + } + proof = brokerTestControl(t, server.URL, brokerRetirePath, c) + if !proof.RetirementProven || !brokerDigestValid(proof.ProofDigest) { + t.Fatal("cancelled prompt's acknowledged session did not retire") + } + creates, inferences, stops, deletes := f.counts() + if creates != 1 || inferences != 0 || stops != 0 || deletes != 1 { + t.Fatalf("cancelled creation replayed work or skipped cleanup: %d %d %d %d", creates, inferences, stops, deletes) + } +} + +func TestBrokerCreationAdmissionRejectionVersusUnknownFailure(t *testing.T) { + for _, mode := range []string{"create-rejected", "create-rejected-server", "create-rejected-conflict", "create-rejected-truncated", "create-rejected-oversized"} { + t.Run(mode, func(t *testing.T) { + f := newBrokerFixture(t, mode) + cfg := brokerTestConfig(t, f) + _, server := startBrokerTest(t, cfg) + c := brokerTestContext(cfg) + status, _, err := brokerTestHTTP(context.Background(), server.URL, brokerResponsesPath, c, brokerTestBody("")) + if err != nil || status == http.StatusOK { + t.Fatal("failed creation exposed an inference result") + } + if mode == "create-rejected" { + proof := brokerTestControl(t, server.URL, brokerSettlePath, c) + if !proof.SettlementProven || proof.RemoteSessionCreated || proof.CreatePending { + t.Fatal("complete creation rejection did not prove no remote session") + } + proof = brokerTestControl(t, server.URL, brokerRetirePath, c) + if !proof.RetirementProven || proof.RemoteSessionCreated { + t.Fatal("rejected creation could not retire without a remote session") + } + } else { + brokerPendingControl(t, server.URL, brokerRetirePath, c, true, 0) + } + creates, inferences, stops, deletes := f.counts() + if creates != 1 || inferences != 0 || stops != 0 || deletes != 0 { + t.Fatal("creation failure was replayed or deleted without ownership proof") + } + }) + } +} + +func TestBrokerRestartSettlesOnlyAcknowledgedInference(t *testing.T) { + for _, acknowledged := range []bool{true, false} { + name, mode := "unknown", "hold-unknown" + if acknowledged { + name, mode = "acknowledged", "hold-known" + } + t.Run(name, func(t *testing.T) { + f := newBrokerFixture(t, mode) + cfg := brokerTestConfig(t, f) + b, server := startBrokerTest(t, cfg) + c := brokerTestContext(cfg) + done := brokerAsyncInference(context.Background(), server.URL, c, brokerTestBody("")) + if acknowledged { + brokerAwait(t, func() bool { return brokerInvocationState(b, c) == "accepted" }) + } else { + select { + case <-f.started: + case <-time.After(3 * time.Second): + t.Fatal("inference did not start") + } + } + b.close() + _ = brokerWaitInference(t, done) + server.Close() + _, server = startBrokerTest(t, cfg) + if acknowledged { + proof := brokerTestControl(t, server.URL, brokerSettlePath, c) + if !proof.SettlementProven { + t.Fatal("original acknowledged owner did not settle after restart") + } + _ = brokerTestControl(t, server.URL, brokerRetirePath, c) + } else { + for range 3 { + brokerPendingControl(t, server.URL, brokerSettlePath, c, false, 1) + brokerPendingControl(t, server.URL, brokerRetirePath, c, false, 1) + } + } + creates, inferences, _, deletes := f.counts() + if creates != 1 || inferences != 1 || (!acknowledged && deletes != 0) || (acknowledged && deletes != 1) { + t.Fatal("restart replayed inference or fabricated deletion proof") + } + }) + } +} + +func TestBrokerRenewalExtendsActiveRequestAndOldLeaseCanClose(t *testing.T) { + f := newBrokerFixture(t, "hold-known") + cfg := brokerTestConfig(t, f) + b, server := startBrokerTest(t, cfg) + c := brokerTestContext(cfg) + originalExpiry := time.Now().Add(800 * time.Millisecond) + c.LeaseExpiresAt = originalExpiry.UTC().Format(time.RFC3339Nano) + done := brokerAsyncInference(context.Background(), server.URL, c, brokerTestBody("")) + brokerAwait(t, func() bool { return brokerInvocationState(b, c) == "accepted" }) + renewal := c + renewal.LeaseGeneration = 2 + renewal.LeaseExpiresAt = time.Now().Add(3 * time.Second).UTC().Format(time.RFC3339Nano) + proof := brokerTestControl(t, server.URL, brokerRenewPath, renewal) + if proof.LeaseGeneration != 2 || proof.State != "open" { + t.Fatal("renewal did not bind the active request") + } + time.Sleep(time.Until(originalExpiry.Add(150 * time.Millisecond))) + _, inferences, stops, _ := f.counts() + if inferences != 1 || stops != 0 || brokerInvocationState(b, c) != "accepted" { + t.Fatal("old expiry canceled a renewed request") + } + proof = brokerTestControl(t, server.URL, brokerSettlePath, c) + if !proof.SettlementProven || proof.LeaseGeneration != 2 { + t.Fatal("old lease could not close exact prompt authority") + } + result := brokerWaitInference(t, done) + if result.status == http.StatusOK { + t.Fatal("settled cancellation exposed output") + } + _ = brokerTestControl(t, server.URL, brokerRetirePath, c) +} + +func TestBrokerKnownAbsentSessionStillRequiresDeleteAcknowledgement(t *testing.T) { + f := newBrokerFixture(t, "success") + cfg := brokerTestConfig(t, f) + _, server := startBrokerTest(t, cfg) + c := brokerTestContext(cfg) + status, _, err := brokerTestHTTP(context.Background(), server.URL, brokerResponsesPath, c, brokerTestBody("")) + if err != nil || status != http.StatusOK { + t.Fatal("initial request failed") + } + // Model DELETE succeeding immediately before a broker loses its response. + f.mu.Lock() + for id := range f.sessions { + delete(f.sessions, id) + } + f.mu.Unlock() + proof := brokerTestControl(t, server.URL, brokerRetirePath, c) + _, _, _, deletes := f.counts() + if !proof.RetirementProven || deletes != 1 { + t.Fatal("known missing target did not obtain DELETE204 + GET404 proof") + } +} + +func TestBrokerCompleteAdmissionRejectionVersusUnknownFailure(t *testing.T) { + for _, mode := range []string{"rejected", "rejected-server", "rejected-truncated", "rejected-oversized"} { + t.Run(mode, func(t *testing.T) { + f := newBrokerFixture(t, mode) + cfg := brokerTestConfig(t, f) + _, server := startBrokerTest(t, cfg) + c := brokerTestContext(cfg) + status, _, err := brokerTestHTTP(context.Background(), server.URL, brokerResponsesPath, c, brokerTestBody("")) + if err != nil || status == http.StatusOK { + t.Fatal("provider rejection exposed a result") + } + if mode == "rejected" { + _ = brokerTestControl(t, server.URL, brokerSettlePath, c) + _ = brokerTestControl(t, server.URL, brokerRetirePath, c) + } else { + brokerPendingControl(t, server.URL, brokerRetirePath, c, false, 1) + } + creates, inferences, _, deletes := f.counts() + if creates != 1 || inferences != 1 || (mode != "rejected" && deletes != 0) { + t.Fatal("uncertain failure was replayed or retired") + } + }) + } +} diff --git a/broker_remote.go b/broker_remote.go new file mode 100644 index 0000000..7ccd19b --- /dev/null +++ b/broker_remote.go @@ -0,0 +1,263 @@ +package main + +import ( + "bytes" + "context" + "encoding/base64" + "encoding/json" + "errors" + "io" + "net" + "net/http" + "net/url" + "strings" + "time" + + "golang.org/x/net/http/httpguts" +) + +var errBrokerRequestUnsent = errors.New("Foundry broker request was not sent") + +type brokerRemoteSession struct { + ID string `json:"agent_session_id"` + Version struct { + Type string `json:"type"` + Version string `json:"agent_version"` + } `json:"version_indicator"` + Status string `json:"status"` +} + +// Complete fallible authentication and request preparation before a caller +// records possible submission. Only sendRemoteRequest crosses that boundary. +func (b *lifecycleBroker) prepareRemoteRequest(ctx context.Context, method, suffix string, body []byte) (*http.Request, error) { + token, err := b.tokenProvider.AccessToken(ctx) + if err != nil { + return nil, errors.Join(errBrokerRequestUnsent, errBrokerRemote) + } + if err := b.pinPrincipal(token); err != nil { + return nil, errors.Join(errBrokerRequestUnsent, err) + } + endpoint := strings.TrimRight(b.cfg.agent.HostedTarget.ProjectEndpoint, "/") + "/agents/" + + url.PathEscape(b.cfg.agent.HostedTarget.AgentName) + suffix + "?api-version=v1" + request, err := http.NewRequestWithContext(ctx, method, endpoint, bytes.NewReader(body)) + if err != nil { + return nil, errors.Join(errBrokerRequestUnsent, errBrokerRemote) + } + // Disable replayable bodies and redirects. A network error never authorizes + // resubmission of a create or inference intent. + request.GetBody = nil + request.Header.Set("Authorization", "Bearer "+token) + request.Header.Set("Foundry-Features", "HostedAgents=V1Preview") + request.Header.Set("Content-Type", "application/json") + request.Header.Set("Accept", "application/json, text/event-stream") + if ctx.Err() != nil { + return nil, errors.Join(errBrokerRequestUnsent, errBrokerClosed) + } + return request, nil +} + +func (b *lifecycleBroker) sendRemoteRequest(request *http.Request) (*http.Response, error) { + if request.Context().Err() != nil { + // Only skipping Do proves non-submission. A cancellation observed + // after dispatch may follow accepted work and remains ambiguous. + if request.Body != nil { + _ = request.Body.Close() + } + return nil, errors.Join(errBrokerRequestUnsent, errBrokerClosed) + } + response, err := b.httpClient.Do(request) + if err != nil { + return nil, errBrokerAmbiguous + } + return response, nil +} + +func newBrokerHTTPClient() *http.Client { + dialer := &net.Dialer{Timeout: 10 * time.Second, KeepAlive: 30 * time.Second} + // Session provisioning has a 45-second operation deadline. Do not abandon + // its one acknowledgement earlier; shorter request contexts still win. + return &http.Client{Transport: &http.Transport{Proxy: nil, DialContext: dialer.DialContext, + TLSHandshakeTimeout: 10 * time.Second, ResponseHeaderTimeout: 45 * time.Second, + MaxIdleConns: 16, MaxIdleConnsPerHost: 8, IdleConnTimeout: 30 * time.Second}, + CheckRedirect: func(*http.Request, []*http.Request) error { return errBrokerRemote }} +} + +func (b *lifecycleBroker) pinPrincipal(token string) error { + if !httpguts.ValidHeaderFieldValue(token) { + return errBrokerRemote + } + parts := strings.Split(token, ".") + if len(parts) != 3 { + return errBrokerRemote + } + raw, err := base64.RawURLEncoding.DecodeString(parts[1]) + var claims struct { + Audience string `json:"aud"` + Tenant string `json:"tid"` + Object string `json:"oid"` + App string `json:"appid"` + AZP string `json:"azp"` + } + if err != nil || acpDecodeStruct(raw, &claims, false) != nil || strings.TrimRight(claims.Audience, "/") != "https://ai.azure.com" || + !acpSafeString(claims.Tenant, 512) || !acpSafeString(claims.Object, 512) { + return errBrokerRemote + } + digest := brokerJSONDigest(claims) + b.mu.Lock() + defer b.mu.Unlock() + if b.ledger.PrincipalDigest == digest { + return b.storageError + } + if b.ledger.PrincipalDigest != "" { + return errBrokerConflict + } + return b.commitLocked(func(next *brokerLedger) error { next.PrincipalDigest = digest; return nil }) +} + +func (b *lifecycleBroker) remoteJSON(ctx context.Context, method, suffix string, body []byte, target any) (int, error) { + request, err := b.prepareRemoteRequest(ctx, method, suffix, body) + if err != nil { + return 0, err + } + return b.remoteJSONRequest(request, target) +} + +func (b *lifecycleBroker) remoteJSONRequest(request *http.Request, target any) (int, error) { + response, err := b.sendRemoteRequest(request) + if err != nil { + return 0, err + } + defer response.Body.Close() //nolint:errcheck + data, err := io.ReadAll(io.LimitReader(response.Body, acpMaxConfigBytes+1)) + if err != nil || len(data) > acpMaxConfigBytes { + return response.StatusCode, errBrokerRemote + } + if target != nil && response.StatusCode >= 200 && response.StatusCode < 300 { + if acpDecodeStruct(data, target, false) != nil { + return response.StatusCode, errBrokerRemote + } + } + return response.StatusCode, nil +} + +func (b *lifecycleBroker) validateRemoteTarget(ctx context.Context) error { + var agent struct { + Name string `json:"name"` + Endpoint struct { + Schemes []json.RawMessage `json:"authorization_schemes"` + } `json:"agent_endpoint"` + } + status, err := b.remoteJSON(ctx, http.MethodGet, "", nil, &agent) + if err != nil || status != http.StatusOK || agent.Name != b.cfg.agent.HostedTarget.AgentName || len(agent.Endpoint.Schemes) == 0 { + return errBrokerRemote + } + // Entra isolation is the verified first increment. Unknown/header isolation + // is not silently adopted from the endpoint or inferred from a broad list. + for _, scheme := range agent.Endpoint.Schemes { + var value struct { + Type string `json:"type"` + } + if acpDecodeStruct(scheme, &value, true) != nil || !strings.EqualFold(value.Type, "entra") { + return errBrokerRemote + } + } + var version struct { + Name string `json:"name"` + Version string `json:"version"` + Status string `json:"status"` + Definition struct { + Kind string `json:"kind"` + } `json:"definition"` + } + status, err = b.remoteJSON(ctx, http.MethodGet, "/versions/"+url.PathEscape(b.cfg.agent.HostedTarget.AgentVersion), nil, &version) + if err != nil || status != http.StatusOK || version.Name != b.cfg.agent.HostedTarget.AgentName || + version.Version != b.cfg.agent.HostedTarget.AgentVersion || version.Status != "active" || version.Definition.Kind != "hosted" { + return errBrokerRemote + } + return nil +} + +func (b *lifecycleBroker) remoteSessionMatches(value brokerRemoteSession, id string) bool { + return value.ID == id && value.Version.Type == "version_ref" && value.Version.Version == b.cfg.agent.HostedTarget.AgentVersion +} + +func brokerSessionSuffix(id string) string { return "/endpoint/sessions/" + url.PathEscape(id) } + +func (b *lifecycleBroker) remoteSessionGet(ctx context.Context, id string) (brokerRemoteSession, int, error) { + var session brokerRemoteSession + status, err := b.remoteJSON(ctx, http.MethodGet, brokerSessionSuffix(id), nil, &session) + if err == nil && status == http.StatusOK && !b.remoteSessionMatches(session, id) { + err = errBrokerConflict + } + return session, status, err +} + +func (b *lifecycleBroker) prepareRemoteSessionCreate(ctx context.Context, id string) (*http.Request, error) { + body, _ := json.Marshal(map[string]any{"agent_session_id": id, + "version_indicator": map[string]string{"type": "version_ref", "agent_version": b.cfg.agent.HostedTarget.AgentVersion}}) + return b.prepareRemoteRequest(ctx, http.MethodPost, "/endpoint/sessions", body) +} + +// A false, nil result proves complete admission rejection. Every possibly sent, +// unacknowledged outcome retains the durable creation intent; an absent-session +// observation cannot clear it. Authentication was completed before the intent. +func (b *lifecycleBroker) remoteSessionCreate(request *http.Request, id string) (bool, error) { + var session brokerRemoteSession + status, err := b.remoteJSONRequest(request, &session) + if errors.Is(err, errBrokerRequestUnsent) { + return false, err + } + if err == nil && brokerDefiniteRejection(status) { + return false, nil + } + if err != nil || status != http.StatusCreated || !b.remoteSessionMatches(session, id) { + return false, errBrokerAmbiguous + } + return true, nil +} + +func (b *lifecycleBroker) remoteSessionStop(ctx context.Context, id string) error { + current, status, err := b.remoteSessionGet(ctx, id) + if err != nil || status != http.StatusOK { + return errBrokerPending + } + status, err = b.remoteJSON(ctx, http.MethodPost, brokerSessionSuffix(id)+":stop", nil, nil) + if err != nil || (status != http.StatusNoContent && status != http.StatusConflict) { + return errBrokerPending + } + // The deployed API returned409 for a duplicate stop. Neither204 nor409 is + // sufficient alone: use exact authenticated identity/version + idle status. + for { + current, status, err = b.remoteSessionGet(ctx, id) + if err != nil || status != http.StatusOK { + return errBrokerPending + } + if current.Status == "idle" { + return nil + } + select { + case <-ctx.Done(): + return errBrokerPending + case <-time.After(100 * time.Millisecond): + } + } +} + +func (b *lifecycleBroker) remoteSessionDelete(ctx context.Context, id string) error { + // Known ownership permits retrying a lost DELETE acknowledgement. An absent + // record does not skip DELETE204 + GET404, and this path is never called for + // an unresolved create or inference intent. + _, status, err := b.remoteSessionGet(ctx, id) + if err != nil || (status != http.StatusOK && status != http.StatusNotFound) { + return errBrokerPending + } + status, err = b.remoteJSON(ctx, http.MethodDelete, brokerSessionSuffix(id), nil, nil) + if err != nil || status != http.StatusNoContent { + return errBrokerPending + } + _, status, err = b.remoteSessionGet(ctx, id) + if err != nil || status != http.StatusNotFound { + return errBrokerPending + } + return nil +} diff --git a/broker_renewal_create_test.go b/broker_renewal_create_test.go new file mode 100644 index 0000000..6ed5b3a --- /dev/null +++ b/broker_renewal_create_test.go @@ -0,0 +1,219 @@ +package main + +import ( + "context" + "encoding/json" + "net/http" + "sync" + "testing" + "time" +) + +func requirePendingCreateRenewal(t *testing.T, proof brokerControlResponse, renewal brokerContext) { + t.Helper() + expires, err := time.Parse(time.RFC3339Nano, proof.LeaseExpiresAt) + wantExpires, _ := time.Parse(time.RFC3339Nano, renewal.LeaseExpiresAt) + if err != nil || proof.State != "open" || proof.LeaseGeneration != renewal.LeaseGeneration || !expires.Equal(wantExpires) || + !proof.CreatePending || proof.ActiveInvocations != 1 || proof.AmbiguousInvocations != 0 || + proof.SettlementProven || proof.RetirementProven || proof.RemoteSessionCreated || proof.ProofDigest != "" { + t.Fatalf("live creation renewal lacks exact acknowledgement: state=%s generation=%d pending=%v active=%d ambiguous=%d", + proof.State, proof.LeaseGeneration, proof.CreatePending, proof.ActiveInvocations, proof.AmbiguousInvocations) + } +} + +func requireBlockedRenewalReplay(t *testing.T, base string, renewal brokerContext, pending bool, ambiguous uint32) { + t.Helper() + proof := brokerTestControl(t, base, brokerRenewPath, renewal) + if proof.State != "blocked" || proof.CreatePending != pending || proof.AmbiguousInvocations != ambiguous || + proof.SettlementProven || proof.RetirementProven || proof.ProofDigest != "" { + t.Fatal("renewal replay reopened unresolved ownership") + } +} + +func TestBrokerRenewalDuringPendingCreateSurvivesOriginalExpiry(t *testing.T) { + f := newBrokerFixture(t, "hold-create") + defer f.unblock() + cfg := brokerTestConfig(t, f) + cfg.operationTimeout = 5 * time.Second + b, server := startBrokerTest(t, cfg) + c := brokerTestContext(cfg) + originalExpiry := time.Now().Add(time.Second) + c.LeaseExpiresAt = originalExpiry.UTC().Format(time.RFC3339Nano) + done := brokerAsyncInference(context.Background(), server.URL, c, brokerTestBody("")) + select { + case <-f.started: + case <-time.After(3 * time.Second): + t.Fatal("original session creation did not reach the held acknowledgement") + } + renewal := c + renewal.LeaseGeneration = 2 + renewal.LeaseExpiresAt = originalExpiry.Add(3 * time.Second).UTC().Format(time.RFC3339Nano) + proof := brokerTestControl(t, server.URL, brokerRenewPath, renewal) + requirePendingCreateRenewal(t, proof, renewal) + // Only this exact idempotent control may be repeated. The original create + // and inference each retain their one attempt. + requirePendingCreateRenewal(t, brokerTestControl(t, server.URL, brokerRenewPath, renewal), renewal) + status := brokerTestControl(t, server.URL, brokerStatusPath, c) + if status.State != "blocked" || !status.CreatePending || status.SettlementProven || status.RetirementProven { + t.Fatal("renewal changed ordinary pending-creation status") + } + time.Sleep(time.Until(originalExpiry.Add(150 * time.Millisecond))) + b.mu.Lock() + prompt := b.ledger.Sessions[brokerJSONDigest(c.Owner)].Prompts[c.promptKey()] + active := !prompt.Closing && !prompt.Settled && prompt.LeaseGeneration == 2 && prompt.LeaseExpiresAt.After(time.Now()) + b.mu.Unlock() + if !active { + t.Fatal("the original expiry closed the renewed creation") + } + creates, inferences, stops, deletes := f.counts() + if creates != 1 || inferences != 0 || stops != 0 || deletes != 0 { + t.Fatal("renewal replayed creation or started inference before creation acknowledgement") + } + f.unblock() + result := brokerWaitInference(t, done) + if result.err != nil || result.status != http.StatusOK { + t.Fatal("the original invocation failed after its renewed creation completed") + } + proof = brokerTestControl(t, server.URL, brokerSettlePath, c) + if !proof.SettlementProven || proof.CreatePending || proof.LeaseGeneration != 2 { + t.Fatal("the original prompt did not settle under its renewed lease") + } + proof = brokerTestControl(t, server.URL, brokerRetirePath, c) + creates, inferences, stops, deletes = f.counts() + if !proof.RetirementProven || creates != 1 || inferences != 1 || stops != 0 || deletes != 1 { + t.Fatal("renewal changed exact-attempt execution or retirement") + } +} + +func TestBrokerRenewalDuringPendingCreateCannotReopenClosing(t *testing.T) { + for _, path := range []string{brokerSettlePath, brokerRetirePath} { + t.Run(path, func(t *testing.T) { + f := newBrokerFixture(t, "hold-create") + defer f.unblock() + cfg := brokerTestConfig(t, f) + _, server := startBrokerTest(t, cfg) + c := brokerTestContext(cfg) + done := brokerAsyncInference(context.Background(), server.URL, c, brokerTestBody("")) + select { + case <-f.started: + case <-time.After(3 * time.Second): + t.Fatal("creation acknowledgement was not held") + } + renewal := c + renewal.LeaseGeneration = 2 + renewal.LeaseExpiresAt = time.Now().Add(15 * time.Second).UTC().Format(time.RFC3339Nano) + requirePendingCreateRenewal(t, brokerTestControl(t, server.URL, brokerRenewPath, renewal), renewal) + brokerPendingControl(t, server.URL, path, c, true, 0) + if path == brokerSettlePath { + requireBlockedRenewalReplay(t, server.URL, renewal, true, 0) + } else { + cc, body := brokerTestControlContext(brokerRenewPath, renewal) + status, _, err := brokerTestHTTP(context.Background(), server.URL, brokerRenewPath, cc, body) + if err != nil || status != http.StatusGone { + t.Fatal("retiring owner accepted a renewal replay") + } + } + f.unblock() + result := brokerWaitInference(t, done) + if result.err == nil && result.status == http.StatusOK { + t.Fatal("closed creation proceeded to inference") + } + proof := brokerTestControl(t, server.URL, brokerRetirePath, c) + creates, inferences, stops, deletes := f.counts() + if !proof.RetirementProven || creates != 1 || inferences != 0 || stops != 0 || deletes != 1 { + t.Fatal("pending-creation cleanup lost ownership or replayed work") + } + }) + } +} + +func TestBrokerRenewalDuringPendingCreateCannotReopenLostAcknowledgement(t *testing.T) { + f := newBrokerFixture(t, "late-create") + defer f.unblock() + createEntered, releaseCreate := make(chan struct{}), make(chan struct{}) + unblock := sync.OnceFunc(func() { close(releaseCreate) }) + defer unblock() + f.createCheck = func(string) { + close(createEntered) + <-releaseCreate + } + cfg := brokerTestConfig(t, f) + b, server := startBrokerTest(t, cfg) + c := brokerTestContext(cfg) + done := brokerAsyncInference(context.Background(), server.URL, c, brokerTestBody("")) + select { + case <-createEntered: + case <-time.After(3 * time.Second): + t.Fatal("creation did not enter the original attempt") + } + renewal := c + renewal.LeaseGeneration = 2 + renewal.LeaseExpiresAt = time.Now().Add(15 * time.Second).UTC().Format(time.RFC3339Nano) + requirePendingCreateRenewal(t, brokerTestControl(t, server.URL, brokerRenewPath, renewal), renewal) + unblock() + result := brokerWaitInference(t, done) + if result.err == nil && result.status == http.StatusOK { + t.Fatal("lost creation acknowledgement exposed inference output") + } + brokerAwait(t, func() bool { return brokerInvocationState(b, c) == "rejected" }) + requireBlockedRenewalReplay(t, server.URL, renewal, true, 0) + b.close() + server.Close() + _, server = startBrokerTest(t, cfg) + requireBlockedRenewalReplay(t, server.URL, renewal, true, 0) + brokerPendingControl(t, server.URL, brokerRetirePath, c, true, 0) + creates, inferences, stops, deletes := f.counts() + if creates != 1 || inferences != 0 || stops != 0 || deletes != 0 { + t.Fatal("abandoned creation was replayed or retired from missing-session evidence") + } + f.unblock() + brokerAwait(t, func() bool { + f.mu.Lock() + defer f.mu.Unlock() + return len(f.sessions) == 1 + }) + // Renewal and a later object cannot replace the lost original CREATE ack. + brokerPendingControl(t, server.URL, brokerRetirePath, c, true, 0) + creates, inferences, stops, deletes = f.counts() + if creates != 1 || inferences != 0 || stops != 0 || deletes != 0 { + t.Fatal("late unacknowledged creation was replayed or retired") + } +} + +func TestBrokerRenewalReplayCannotClearAmbiguousInference(t *testing.T) { + f := newBrokerFixture(t, "hold-unknown") + defer f.unblock() + cfg := brokerTestConfig(t, f) + b, server := startBrokerTest(t, cfg) + c := brokerTestContext(cfg) + ctx, cancel := context.WithCancel(context.Background()) + defer cancel() + done := brokerAsyncInference(ctx, server.URL, c, brokerTestBody("")) + select { + case <-f.started: + case <-time.After(3 * time.Second): + t.Fatal("inference did not enter the original request") + } + renewal := c + renewal.LeaseGeneration = 2 + renewal.LeaseExpiresAt = time.Now().Add(15 * time.Second).UTC().Format(time.RFC3339Nano) + proof := brokerTestControl(t, server.URL, brokerRenewPath, renewal) + if proof.State != "open" || proof.LeaseGeneration != 2 { + t.Fatal("live inference did not acknowledge its lease renewal") + } + cancel() + _ = brokerWaitInference(t, done) + brokerAwait(t, func() bool { return brokerInvocationState(b, c) == "uncertain" }) + requireBlockedRenewalReplay(t, server.URL, renewal, false, 1) + brokerPendingControl(t, server.URL, brokerRetirePath, c, false, 1) + creates, inferences, _, deletes := f.counts() + if creates != 1 || inferences != 1 || deletes != 0 { + t.Fatal("ambiguous inference was replayed or deleted") + } + statusContext, body := brokerTestControlContext(brokerStatusPath, c) + status, data, err := brokerTestHTTP(context.Background(), server.URL, brokerStatusPath, statusContext, body) + if err != nil || status != http.StatusOK || json.Unmarshal(data, &proof) != nil || proof.State != "blocked" || + proof.SettlementProven || proof.RetirementProven { + t.Fatal("ambiguous ownership no longer blocks status and cleanup") + } +} diff --git a/broker_response_identity_write_test.go b/broker_response_identity_write_test.go new file mode 100644 index 0000000..4fac3ce --- /dev/null +++ b/broker_response_identity_write_test.go @@ -0,0 +1,295 @@ +package main + +import ( + "bytes" + "context" + "encoding/json" + "errors" + "io" + "os" + "path/filepath" + "strings" + "sync" + "testing" + "time" +) + +const brokerIdentityRemoteSession = "11111111-1111-4111-8111-111111111111" + +func newBrokerResponseIdentityFixture(t *testing.T) (*lifecycleBroker, brokerContext) { + t.Helper() + cfg := brokerConfiguration{configDigest: brokerSHA([]byte("identity-write-fixture")), stateDir: filepath.Join(t.TempDir(), "broker")} + store, ledger, err := openBrokerStore(cfg.stateDir, cfg.configDigest) + if err != nil { + t.Fatal("could not initialize identity fixture store") + } + t.Cleanup(store.close) + ctx, cancel := context.WithCancel(context.Background()) + t.Cleanup(cancel) + b := &lifecycleBroker{store: store, ledger: ledger, ctx: ctx, cancel: cancel, active: map[string]brokerActive{}} + c := brokerTestContext(cfg) + c.BodySHA256 = brokerSHA([]byte("identity-write-input")) + c.LeaseExpiresAt = time.Now().Add(time.Minute).UTC().Format(time.RFC3339Nano) + err = b.commitLocked(func(next *brokerLedger) error { + session, err := brokerEnsureSession(next, c) + if err != nil { + return err + } + next.PrincipalDigest = brokerSHA([]byte("identity-write-principal")) + session.RemoteID, session.CreateState = brokerIdentityRemoteSession, "known" + if _, err := brokerRecordOperation(session, brokerResponsesPath, c); err != nil { + return err + } + prompt, err := brokerEnsurePrompt(session, c) + if err != nil { + return err + } + prompt.LastSequence = c.InvocationSequence + prompt.Invocations[c.InvocationSequence] = &brokerInvocation{Sequence: c.InvocationSequence, + OperationID: c.OperationID, BodyDigest: c.BodySHA256, State: "intent"} + return nil + }) + if err != nil || !brokerLedgerValid(b.ledger, cfg.configDigest) { + t.Fatal("identity fixture lacks valid durable invocation intent") + } + return b, c +} + +// The store atomically replaces state.json on every save. Supplying exactly one +// SSE frame per Read lets the next Read count that replacement after its flush, +// without replacing the production store or relying on timestamp resolution. +type brokerIdentityEventReader struct { + t *testing.T + path string + events []string + index int + last os.FileInfo + writes int +} + +func (r *brokerIdentityEventReader) observe() { + r.t.Helper() + info, err := os.Stat(r.path) + if err != nil { + r.t.Fatal("durable identity record unavailable") + } + if r.last != nil && !os.SameFile(r.last, info) { + r.writes++ + } + r.last = info +} + +func (r *brokerIdentityEventReader) Read(p []byte) (int, error) { + r.observe() + if r.index == 1 && r.writes != 1 { + r.t.Fatal("first acceptance was not persisted before reading the next event") + } + if r.index == len(r.events) { + return 0, io.EOF + } + event := r.events[r.index] + if len(p) < len(event) { + return 0, io.ErrShortBuffer + } + r.index++ + return copy(p, event), nil +} + +func brokerIdentityLedgerBytes(t *testing.T, b *lifecycleBroker) []byte { + t.Helper() + data, err := os.ReadFile(filepath.Join(b.store.dir, "state.json")) + if err != nil { + t.Fatal("could not read durable identity fixture") + } + var ledger brokerLedger + if json.Unmarshal(data, &ledger) != nil || !brokerLedgerValid(&ledger, b.ledger.ConfigDigest) || + brokerJSONDigest(ledger) != brokerJSONDigest(b.ledger) { + t.Fatal("in-memory identity does not match valid durable ownership") + } + return data +} + +func TestBrokerResponseIdentityWritesOncePerInvocation(t *testing.T) { + for _, test := range []struct { + name string + events int + valid bool + }{ + {"repeated", 64, true}, + {"event-limit", defaultMaxEvents, true}, + {"over-event-limit", defaultMaxEvents + 1, false}, + } { + t.Run(test.name, func(t *testing.T) { + b, c := newBrokerResponseIdentityFixture(t) + events := make([]string, test.events) + for i := range events { + events[i] = acpTestSSE(`{"type":"response.in_progress","response":{"id":"response-1","status":"in_progress"}}`) + } + events[0] = acpTestSSE(`{"type":"response.created","response":{"id":"response-1","status":"in_progress"}}`) + events[len(events)-1] = acpTestSSE(`{"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) + if err != nil { + t.Fatal("coherent acceptance evidence was rejected") + } + if reader.writes != 1 { + t.Fatalf("identity persistence count = %d for %d lifecycle events; want 1", reader.writes, test.events) + } + invocation := b.ledger.Sessions[brokerJSONDigest(c.Owner)].Prompts[c.promptKey()].Invocations[c.InvocationSequence] + if invocation.State != "accepted" || invocation.ResponseID != "response-1" || invocation.ResponseAlias == "" { + t.Fatal("stream tracking did not preserve acceptance separately from completion") + } + _ = brokerIdentityLedgerBytes(t, b) + summary, err := acpParseFoundrySSE(bytes.NewReader(data)) + if (err == nil) != test.valid { + t.Fatal("identity deduplication changed stream event validation") + } + if test.valid { + if _, err := b.commitCompletedResponse(c, summary); err != nil { + t.Fatal("validated response could not persist completion") + } + reader.observe() + if reader.writes != 2 { + t.Fatal("completion did not perform its separate durable write") + } + _ = brokerIdentityLedgerBytes(t, b) + } + }) + } +} + +func TestBrokerResponseIdentityDuplicateStillRejectsConflicts(t *testing.T) { + b, c := newBrokerResponseIdentityFixture(t) + accepted := foundryResponse{ID: "response-1", AgentSessionID: brokerIdentityRemoteSession} + if err := b.recordResponseIdentity(c, accepted, brokerIdentityRemoteSession); err != nil { + t.Fatal("initial response identity was rejected") + } + before := brokerIdentityLedgerBytes(t, b) + for name, response := range map[string]foundryResponse{ + "changed-response": {ID: "response-2"}, + "missing-response": {}, + "invalid-response": {ID: "response/invalid"}, + "changed-session": {ID: accepted.ID, AgentSessionID: "another-session"}, + } { + t.Run(name, func(t *testing.T) { + if err := b.recordResponseIdentity(c, response, brokerIdentityRemoteSession); !errors.Is(err, errBrokerConflict) { + t.Fatal("conflicting response identity was accepted") + } + if !bytes.Equal(before, brokerIdentityLedgerBytes(t, b)) { + t.Fatal("conflicting identity changed durable ownership") + } + }) + } + if _, err := b.commitCompletedResponse(c, foundryStreamSummary{ResponseID: accepted.ID}); err != nil { + t.Fatal("could not prepare previous response identity") + } + nextContext := c + nextContext.InvocationSequence++ + nextContext.OperationID = "identity-next-invocation" + if err := b.commitLocked(func(next *brokerLedger) error { + session := next.Sessions[brokerJSONDigest(c.Owner)] + if _, err := brokerRecordOperation(session, brokerResponsesPath, nextContext); err != nil { + return err + } + prompt := session.Prompts[c.promptKey()] + prompt.LastSequence = nextContext.InvocationSequence + prompt.Invocations[nextContext.InvocationSequence] = &brokerInvocation{Sequence: nextContext.InvocationSequence, + OperationID: nextContext.OperationID, BodyDigest: nextContext.BodySHA256, State: "intent"} + return nil + }); err != nil { + t.Fatal("could not prepare next invocation") + } + before = brokerIdentityLedgerBytes(t, b) + if err := b.recordResponseIdentity(nextContext, accepted, brokerIdentityRemoteSession); !errors.Is(err, errBrokerConflict) { + t.Fatal("response identity from an earlier invocation was reused") + } + if !bytes.Equal(before, brokerIdentityLedgerBytes(t, b)) { + t.Fatal("reused response changed durable ownership") + } +} + +func TestBrokerResponseIdentityConcurrentDuplicatesDoNotWrite(t *testing.T) { + b, c := newBrokerResponseIdentityFixture(t) + response := foundryResponse{ID: "response-1"} + if err := b.recordResponseIdentity(c, response, brokerIdentityRemoteSession); err != nil { + t.Fatal("initial response identity was rejected") + } + reader := &brokerIdentityEventReader{t: t, path: filepath.Join(b.store.dir, "state.json")} + // Hold the inode so multiple incorrect rewrites cannot reuse it before the + // final observation and hide a persistence call. + record, err := os.Open(reader.path) + if err != nil { + t.Fatal("could not pin the durable identity record") + } + defer func() { _ = record.Close() }() + reader.observe() + before := brokerIdentityLedgerBytes(t, b) + var group sync.WaitGroup + for range 32 { + group.Go(func() { + if err := b.recordResponseIdentity(c, response, brokerIdentityRemoteSession); err != nil { + t.Error("concurrent duplicate identity was rejected") + } + }) + } + group.Wait() + reader.observe() + if reader.writes != 0 || !bytes.Equal(before, brokerIdentityLedgerBytes(t, b)) { + t.Fatal("concurrent duplicate identities rewrote durable ownership") + } +} + +func TestBrokerResponseIdentityPersistenceFailureRemainsClosed(t *testing.T) { + for _, existing := range []bool{false, true} { + name := "first-acceptance" + if existing { + name = "duplicate-after-store-failure" + } + t.Run(name, func(t *testing.T) { + b, c := newBrokerResponseIdentityFixture(t) + response := foundryResponse{ID: "response-1"} + if existing && b.recordResponseIdentity(c, response, brokerIdentityRemoteSession) != nil { + t.Fatal("initial response identity was rejected") + } + before := brokerIdentityLedgerBytes(t, b) + active, cancel := context.WithCancel(context.Background()) + defer cancel() + b.active[brokerJSONDigest(c.Owner)] = brokerActive{prompt: c.promptKey(), cancel: cancel} + if err := os.Rename(b.store.dir, b.store.dir+"-unavailable"); err != nil { + t.Fatal("could not inject identity persistence failure") + } + if existing && !errors.Is(b.commitLocked(func(*brokerLedger) error { return nil }), errBrokerStorage) { + t.Fatal("store failure did not poison the broker") + } + if err := b.recordResponseIdentity(c, response, brokerIdentityRemoteSession); !errors.Is(err, errBrokerStorage) { + t.Fatal("identity callback bypassed failed persistence") + } + if active.Err() != context.Canceled || !errors.Is(b.storageError, errBrokerStorage) { + t.Fatal("persistence failure did not contain active work") + } + if brokerSHA(before) != brokerJSONDigest(b.ledger) { + t.Fatal("failed persistence changed the last durable identity") + } + if err := os.Rename(b.store.dir+"-unavailable", b.store.dir); err != nil { + t.Fatal("could not restore fixture store for readback") + } + if !bytes.Equal(before, brokerIdentityLedgerBytes(t, b)) { + t.Fatal("failed identity write changed the durable record") + } + }) + } +} + +func TestBrokerResponseIdentityTrackingStillRejectsMalformedTail(t *testing.T) { + b, c := newBrokerResponseIdentityFixture(t) + created := acpTestSSE(`{"type":"response.created","response":{"id":"response-1","status":"in_progress"}}`) + stream := created + acpTestSSE(`{"type":"response.in_progress","response":{"id":"response-1","status":"completed"}}`) + if _, err := b.readTrackedStream(strings.NewReader(stream), c, brokerIdentityRemoteSession); err == nil { + t.Fatal("duplicate response identity bypassed lifecycle validation") + } + invocation := b.ledger.Sessions[brokerJSONDigest(c.Owner)].Prompts[c.promptKey()].Invocations[c.InvocationSequence] + if invocation.ResponseID != "response-1" || invocation.State != "accepted" { + t.Fatal("malformed tail erased earlier durable acknowledgement") + } + _ = brokerIdentityLedgerBytes(t, b) +} diff --git a/broker_response_storage_precedence_test.go b/broker_response_storage_precedence_test.go new file mode 100644 index 0000000..0d3fcf9 --- /dev/null +++ b/broker_response_storage_precedence_test.go @@ -0,0 +1,106 @@ +package main + +import ( + "bytes" + "context" + "errors" + "io" + "net/http" + "os" + "path/filepath" + "strings" + "sync" + "testing" + "time" +) + +type brokerGatedResponseBody struct { + io.ReadCloser + onRead, onClose func() +} + +func (b *brokerGatedResponseBody) Read(p []byte) (int, error) { + if b.onRead != nil { + b.onRead() + } + return b.ReadCloser.Read(p) +} + +func (b *brokerGatedResponseBody) Close() error { + if b.onClose != nil { + b.onClose() + } + return b.ReadCloser.Close() +} + +func TestBrokerStoragePoisonResponsePrecedence(t *testing.T) { + for _, phase := range []string{"response-identity", "finish-after-durable-completion"} { + t.Run(phase, func(t *testing.T) { + f := newBrokerFixture(t, "success") + cfg := brokerTestConfig(t, f) + entered, release := make(chan struct{}), make(chan struct{}) + unblock := sync.OnceFunc(func() { close(release) }) + defer unblock() + gate := sync.OnceFunc(func() { close(entered); <-release }) + client := newBrokerHTTPClient() + transport := client.Transport.(*http.Transport) + defer transport.CloseIdleConnections() + client.Transport = brokerFixtureTransport(func(r *http.Request) (*http.Response, error) { + response, err := transport.RoundTrip(r) + if err == nil && r.Method == http.MethodPost && strings.HasSuffix(r.URL.Path, "/protocols/openai/responses") { + body := &brokerGatedResponseBody{ReadCloser: response.Body} + if phase == "response-identity" { + body.onRead = gate + } else { + body.onClose = gate + } + response.Body = body + } + return response, err + }) + b, server := startBrokerTestWithClient(t, cfg, client) + c := brokerTestContext(cfg) + done := brokerAsyncInference(context.Background(), server.URL, c, brokerTestBody("")) + select { + case <-entered: + case <-time.After(3 * time.Second): + t.Fatal("original response did not reach the selected persistence boundary") + } + before, err := os.ReadFile(filepath.Join(cfg.stateDir, "state.json")) + if err != nil { + t.Fatal(err) + } + state := brokerInvocationState(b, c) + wantState := "intent" + if phase == "finish-after-durable-completion" { + wantState = "completed" + } + if state != wantState { + t.Fatalf("wrong selected boundary: state %q, want %q", state, wantState) + } + retained, restore := brokerReviewFaultStore(t, cfg.stateDir) + unblock() + result := brokerWaitInference(t, done) + b.mu.Lock() + poisoned := errors.Is(b.storageError, errBrokerStorage) + cancelled := b.ctx.Err() != nil + active := len(b.active) + b.mu.Unlock() + b.close() + server.Close() + after, err := os.ReadFile(filepath.Join(retained, "state.json")) + restore() + if err != nil || !bytes.Equal(before, after) || !poisoned || !cancelled || active != 0 { + t.Fatal("storage failure did not retain the exact durable owner and remove active authority") + } + creates, inference, stops, deletes := f.counts() + if creates != 1 || inference != 1 || stops != 0 || deletes != 0 { + t.Fatal("storage failure replayed work or claimed cleanup") + } + t.Logf("phase=%s HTTP=%d storagePoisoned=%v originalDurableState=%s cancelled=%v", phase, result.status, poisoned, state, cancelled) + if result.err != nil || result.status != http.StatusServiceUnavailable { + t.Errorf("poisoned response HTTP = %d, want 503", result.status) + } + }) + } +} diff --git a/broker_responses.go b/broker_responses.go new file mode 100644 index 0000000..75ae403 --- /dev/null +++ b/broker_responses.go @@ -0,0 +1,584 @@ +package main + +import ( + "bufio" + "bytes" + "context" + "encoding/json" + "errors" + "io" + "mime" + "net/http" + "strings" + "time" + + "github.com/google/uuid" +) + +func (b *lifecycleBroker) serveResponses(w http.ResponseWriter, r *http.Request, c brokerContext, raw []byte) { + var request acpResponseRequest + if acpDecode(raw, &request, true) != nil || request.Model != b.cfg.agent.Model || !request.Stream || !request.Store || request.Input == nil { + brokerWriteError(w, errBrokerInvalid) + return + } + var fields map[string]json.RawMessage + if json.Unmarshal(raw, &fields) != nil { + brokerWriteError(w, errBrokerInvalid) + return + } + // Match encoding/json's field-name folding, including presence with null + // or empty values that would disappear when the decoded request is encoded. + for name := range fields { + if strings.EqualFold(name, "agent_session_id") || + (b.cfg.agent.ToolSchemaMode == toolSchemaModeProviderStatic && strings.EqualFold(name, "tools")) { + brokerWriteError(w, errBrokerInvalid) + return + } + } + for _, tool := range request.Tools { + var schema map[string]any + if tool.Type != "function" || !acpSafeString(tool.Name, 512) || acpDecode(tool.Parameters, &schema, false) != nil { + brokerWriteError(w, errBrokerInvalid) + return + } + } + key := brokerJSONDigest(c.Owner) + runCtx, cancel := context.WithCancel(b.ctx) + err := b.reserveInvocation(c, &request, cancel) + if err != nil { + cancel() + brokerWriteError(w, err) + return + } + defer b.wg.Done() + defer cancel() + stopDisconnect := context.AfterFunc(r.Context(), func() { b.closePrompt(key, c.promptKey()); cancel() }) + defer stopDisconnect() + var result foundryResponse + result, err = b.invoke(runCtx, c, request.foundryResponseRequest) + stopDisconnect() + b.finishInvocation(c, err) + // Finalization can poison storage after a valid remote completion. The + // durable failure also takes precedence over earlier transport errors. + b.mu.Lock() + if b.storageError != nil { + err = b.storageError + } + b.mu.Unlock() + if err != nil { + brokerWriteError(w, err) + return + } + // Buffer only within the advertised response bound. The ACP child also waits + // for explicit completion before proposing any tool call. Never stream a + // terminal result that has not been durably linked to this owner. + _ = http.NewResponseController(w).SetWriteDeadline(time.Now().Add(5 * time.Second)) + w.Header().Set("Content-Type", "application/json") + w.WriteHeader(http.StatusOK) + if json.NewEncoder(w).Encode(result) != nil { + b.closePrompt(key, c.promptKey()) + } +} + +func (b *lifecycleBroker) reserveInvocation(c brokerContext, request *acpResponseRequest, cancel context.CancelFunc) error { + b.mu.Lock() + defer b.mu.Unlock() + if b.storageError != nil { + return b.storageError + } + if b.ctx.Err() != nil { + return errBrokerClosed + } + key := brokerJSONDigest(c.Owner) + if _, active := b.active[key]; active { + return errBrokerConflict + } + err := b.commitCapacityLocked(true, func(next *brokerLedger) error { + session, err := brokerEnsureSession(next, c) + if err != nil { + return err + } + if session.Retiring || session.Retired { + return errBrokerClosed + } + duplicate, err := brokerRecordOperation(session, brokerResponsesPath, c) + if err != nil { + return err + } + if duplicate { + return errBrokerConflict + } // No response replay cache or inference retry. + prompt, err := brokerEnsurePrompt(session, c) + if err != nil { + return err + } + if prompt.Closing || prompt.Settled || !prompt.LeaseExpiresAt.After(time.Now()) { + return errBrokerClosed + } + expires, _ := time.Parse(time.RFC3339Nano, c.LeaseExpiresAt) + if c.LeaseGeneration != prompt.LeaseGeneration || !expires.Equal(prompt.LeaseExpiresAt) || c.InvocationSequence <= prompt.LastSequence { + return errBrokerConflict + } + if prompt.LastSequence != 0 && request.PreviousResponseID != prompt.LastAlias { + return errBrokerConflict + } + if err := brokerTranslatePrevious(session, c, prompt.LastSequence == 0, &request.foundryResponseRequest); err != nil { + return err + } + prompt.LastSequence = c.InvocationSequence + prompt.Invocations[c.InvocationSequence] = &brokerInvocation{Sequence: c.InvocationSequence, OperationID: c.OperationID, + BodyDigest: c.BodySHA256, State: "reserved"} + return nil + }) + if err != nil { + return err + } + b.active[key] = brokerActive{prompt: c.promptKey(), cancel: cancel} + b.wg.Add(1) + return nil +} + +func brokerTranslatePrevious(session *brokerSession, c brokerContext, first bool, request *foundryResponseRequest) error { + previous, hasPrevious := session.Responses[request.PreviousResponseID] + if request.PreviousResponseID != "" && (!hasPrevious || !previous.Completed || + (first && previous.HasFunctions)) { + return errBrokerConflict + } + if !first && (!hasPrevious || previous.PromptKey != c.promptKey()) { + return errBrokerConflict + } + if first { + if _, ok := request.Input.(string); !ok { + // This bridge sends one user text on the first round. Native item + // references and child-chosen provider identities are not accepted. + return errBrokerInvalid + } + if hasPrevious { + request.PreviousResponseID = previous.RemoteID + } + return nil + } + inputs, list := request.Input.([]any) + if !list || !previous.HasFunctions || len(inputs) != len(previous.CallIDs) { + return errBrokerConflict + } + seen := map[string]bool{} + for _, input := range inputs { + item, object := input.(map[string]any) + if !object || item["type"] != "function_call_output" { + return errBrokerConflict + } + call, ok := item["call_id"].(string) + output, outputOK := item["output"].(string) + remote, owned := previous.CallIDs[call] + if !ok || !outputOK || !owned || seen[call] || len(item) != 3 || len(output) > defaultMaxBrokeredBytes { + return errBrokerConflict + } + seen[call] = true + item["call_id"] = remote + } + request.PreviousResponseID = previous.RemoteID + return nil +} + +func (b *lifecycleBroker) invoke(ctx context.Context, c brokerContext, request foundryResponseRequest) (foundryResponse, error) { + key, promptKey := brokerJSONDigest(c.Owner), c.promptKey() + b.mu.Lock() + session := b.ledger.Sessions[key] + createState, remoteID := session.CreateState, session.RemoteID + b.mu.Unlock() + if createState == "none" { + if err := b.validateRemoteTarget(ctx); err != nil { + return foundryResponse{}, err + } + remoteID = uuid.NewString() + prepareCtx, cancelPrepare := context.WithTimeout(ctx, b.cfg.operationTimeout) + prepared, err := b.prepareRemoteSessionCreate(prepareCtx, remoteID) + deadline, _ := prepareCtx.Deadline() + cancelPrepare() + if err != nil { + return foundryResponse{}, err + } + // Before durable intent, caller cancellation must leave creation unsent. + // After intent, keep this one bounded attempt alive to retain its ack. + createCtx, cancelCreate := context.WithDeadline(b.ctx, deadline) + prepared = prepared.WithContext(createCtx) + b.mu.Lock() + err = b.commitLocked(func(next *brokerLedger) error { + session := next.Sessions[key] + prompt := session.Prompts[promptKey] + if session.Retiring || prompt.Closing || ctx.Err() != nil || createCtx.Err() != nil || !prompt.LeaseExpiresAt.After(time.Now()) { + return errBrokerClosed + } + session.RemoteID, session.CreateState = remoteID, "intent" + return nil + }) + b.mu.Unlock() + if err != nil { + cancelCreate() + return foundryResponse{}, err + } + // Once the intent is durable, finish this one creation attempt even if + // the prompt closes. Losing its acknowledgement would leave an owner + // that cannot be retired from a later 404. Broker shutdown and the + // operation deadline still bound the attempt; it is never replayed. + created, createErr := b.remoteSessionCreate(prepared, remoteID) + cancelCreate() + if createErr != nil && !errors.Is(createErr, errBrokerRequestUnsent) { + return foundryResponse{}, createErr + } + b.mu.Lock() + err = b.commitLocked(func(next *brokerLedger) error { + session := next.Sessions[key] + if created { + session.CreateState = "known" + } else { + // Deliberately skipped dispatch or complete admission rejection + // proves no session was created. Ambiguity never enters here. + session.RemoteID, session.CreateState = "", "none" + } + return nil + }) + b.mu.Unlock() + if err != nil { + return foundryResponse{}, err + } + if createErr != nil { + return foundryResponse{}, createErr + } + if !created { + return foundryResponse{}, errBrokerRemote + } + } else if createState != "known" { + return foundryResponse{}, errBrokerPending + } + if ctx.Err() != nil { + return foundryResponse{}, errBrokerClosed + } + current, status, err := b.remoteSessionGet(ctx, remoteID) + if err != nil || status != http.StatusOK || (current.Status != "active" && current.Status != "idle") { + return foundryResponse{}, errBrokerRemote + } + request.AgentSessionID = remoteID + body, err := json.Marshal(request) + if err != nil { + return foundryResponse{}, errBrokerInvalid + } + prepared, err := b.prepareRemoteRequest(ctx, http.MethodPost, "/endpoint/protocols/openai/responses", body) + if err != nil { + return foundryResponse{}, err + } + b.mu.Lock() + err = b.commitLocked(func(next *brokerLedger) error { + session := next.Sessions[key] + prompt := session.Prompts[promptKey] + if session.Retiring || prompt.Closing || ctx.Err() != nil || !prompt.LeaseExpiresAt.After(time.Now()) { + return errBrokerClosed + } + prompt.Invocations[c.InvocationSequence].State = "intent" + return nil + }) + b.mu.Unlock() + if err != nil { + return foundryResponse{}, err + } + response, err := b.sendRemoteRequest(prepared) + if err != nil { + return foundryResponse{}, err + } + defer response.Body.Close() //nolint:errcheck + if response.StatusCode != http.StatusOK { + count, readErr := io.Copy(io.Discard, io.LimitReader(response.Body, acpMaxConfigBytes+1)) + if readErr != nil || count > acpMaxConfigBytes || !brokerDefiniteRejection(response.StatusCode) { + return foundryResponse{}, errBrokerAmbiguous + } + // An explicit complete HTTP rejection is different from a lost response. + // It never authorizes a retry, but has no delayed unacknowledged request. + b.mu.Lock() + err = b.commitLocked(func(next *brokerLedger) error { + next.Sessions[key].Prompts[promptKey].Invocations[c.InvocationSequence].State = "rejected" + return nil + }) + b.mu.Unlock() + if err != nil { + return foundryResponse{}, err + } + return foundryResponse{}, errBrokerRemote + } + mediaType, _, err := mime.ParseMediaType(response.Header.Get("Content-Type")) + if err != nil { + return foundryResponse{}, errBrokerAmbiguous + } + var summary foundryStreamSummary + switch mediaType { + case "text/event-stream": + var data []byte + data, err = b.readTrackedStream(response.Body, c, remoteID) + if err == nil { + summary, err = acpParseFoundrySSE(bytes.NewReader(data)) + } + case "application/json": + var data []byte + data, err = io.ReadAll(io.LimitReader(response.Body, defaultMaxStreamBytes+1)) + if err == nil && len(data) <= defaultMaxStreamBytes { + var document foundryResponse + document, err = brokerDecodeResponseEvidence(data) + if err == nil { + err = b.recordResponseIdentity(c, document, remoteID) + } + if err == nil && document.Status == "completed" { + document, err = acpDecodeFoundryResponse(data) + if err == nil { + summary, err = processCompletedResponse(document, responseCallbacks{}) + } + } else if err == nil { + err = errBrokerRemote + } + } else { + err = errBrokerAmbiguous + } + default: + err = errBrokerAmbiguous + } + if err != nil || acpValidateSummary(summary) != nil { + return foundryResponse{}, errBrokerAmbiguous + } + return b.commitCompletedResponse(c, summary) +} + +func brokerDefiniteRejection(status int) bool { + // A gateway timeout or server error may follow a forwarded request whose + // response was lost. Only explicit admission rejections close this ambiguity. + switch status { + case http.StatusBadRequest, http.StatusUnauthorized, http.StatusForbidden, + http.StatusNotFound, http.StatusMethodNotAllowed, http.StatusRequestEntityTooLarge, + http.StatusUnsupportedMediaType, http.StatusUnprocessableEntity, http.StatusTooManyRequests: + return true + default: + return false + } +} + +// Acceptance evidence is distinct from usable output. A coherent failed or +// incomplete response still acknowledges this invocation and permits stop +// containment, while only the stricter ACP decoder may admit successful output. +func brokerDecodeResponseEvidence(data []byte) (foundryResponse, error) { + var evidence struct { + ID string `json:"id"` + Status string `json:"status"` + AgentSessionID string `json:"agent_session_id"` + Output json.RawMessage `json:"output"` + Error *foundryError `json:"error"` + Incomplete *foundryIncomplete `json:"incomplete_details"` + } + if acpDecodeStruct(data, &evidence, false) != nil || evidence.ID == "" || validateProviderIdentifier("response", evidence.ID) != nil { + return foundryResponse{}, errBrokerRemote + } + response := foundryResponse{ID: evidence.ID, Status: evidence.Status, AgentSessionID: evidence.AgentSessionID, + Error: evidence.Error, Incomplete: evidence.Incomplete} + switch response.Status { + case "queued", "in_progress", "completed", "cancelled": + if response.Error != nil || response.Incomplete != nil { + return foundryResponse{}, errBrokerRemote + } + case "failed": + if response.Incomplete != nil { + return foundryResponse{}, errBrokerRemote + } + case "incomplete": + if response.Error != nil { + return foundryResponse{}, errBrokerRemote + } + default: + return foundryResponse{}, errBrokerRemote + } + return response, nil +} + +func (b *lifecycleBroker) readTrackedStream(reader io.Reader, c brokerContext, remoteID string) ([]byte, error) { + limited := &io.LimitedReader{R: reader, N: defaultMaxStreamBytes + 1} + scanner := bufio.NewScanner(limited) + scanner.Buffer(make([]byte, 32<<10), defaultMaxEventBytes) + var raw, event []byte + flush := func() error { + if len(event) == 0 || bytes.Equal(bytes.TrimSpace(event), []byte("[DONE]")) { + return nil + } + var frame struct { + Type string `json:"type"` + Response json.RawMessage `json:"response"` + Error *foundryError `json:"error"` + } + if acpDecodeStruct(event, &frame, false) != nil { + return errBrokerRemote + } + if frame.Response == nil { + return nil + } + response, err := brokerDecodeResponseEvidence(frame.Response) + if err != nil || frame.Error != nil { + return errBrokerRemote + } + switch frame.Type { + case "response.created": + if response.Status != "queued" && response.Status != "in_progress" { + return errBrokerRemote + } + case "response.queued", "response.in_progress", "response.completed", "response.failed", "response.incomplete", "response.cancelled": + if frame.Type != "response."+response.Status { + return errBrokerRemote + } + default: + return errBrokerRemote + } + return b.recordResponseIdentity(c, response, remoteID) + } + for scanner.Scan() { + line := scanner.Bytes() + if len(raw)+len(line)+1 > defaultMaxStreamBytes { + return nil, errBrokerRemote + } + raw = append(raw, line...) + raw = append(raw, '\n') + if len(line) == 0 { + if err := flush(); err != nil { + return nil, err + } + event = nil + } else if part, ok := bytes.CutPrefix(line, []byte("data:")); ok { + part = bytes.TrimPrefix(part, []byte(" ")) + if len(event)+len(part)+1 > defaultMaxEventBytes { + return nil, errBrokerRemote + } + if len(event) > 0 { + event = append(event, '\n') + } + event = append(event, part...) + } + } + if scanner.Err() != nil || limited.N <= 0 || len(event) != 0 { + return nil, errBrokerAmbiguous + } + return raw, nil +} + +func (b *lifecycleBroker) recordResponseIdentity(c brokerContext, response foundryResponse, remoteID string) error { + if response.ID == "" || validateProviderIdentifier("response", response.ID) != nil || + (response.AgentSessionID != "" && response.AgentSessionID != remoteID) { + return errBrokerConflict + } + b.mu.Lock() + defer b.mu.Unlock() + if b.storageError != nil { + return b.storageError + } + key := brokerJSONDigest(c.Owner) + // Lifecycle events repeat the same identity. Only its first acceptance + // changes durable ownership, but duplicates must still fail on storage loss. + if previous := b.ledger.Sessions[key].Prompts[c.promptKey()].Invocations[c.InvocationSequence].ResponseID; previous != "" { + if previous != response.ID { + return errBrokerConflict + } + return nil + } + return b.commitLocked(func(next *brokerLedger) error { + session := next.Sessions[key] + invocation := session.Prompts[c.promptKey()].Invocations[c.InvocationSequence] + for _, previous := range session.Responses { + if previous.RemoteID == response.ID { + return errBrokerConflict + } + } + alias := "fr_" + uuid.NewString() + invocation.ResponseID, invocation.ResponseAlias, invocation.State = response.ID, alias, "accepted" + session.Responses[alias] = brokerResponseID{RemoteID: response.ID, PromptKey: c.promptKey()} + return nil + }) +} + +func (b *lifecycleBroker) commitCompletedResponse(c brokerContext, summary foundryStreamSummary) (foundryResponse, error) { + b.mu.Lock() + defer b.mu.Unlock() + var output foundryResponse + err := b.commitCapacityLocked(true, func(next *brokerLedger) error { + session := next.Sessions[brokerJSONDigest(c.Owner)] + prompt := session.Prompts[c.promptKey()] + invocation := prompt.Invocations[c.InvocationSequence] + if invocation.ResponseID != summary.ResponseID || invocation.ResponseAlias == "" { + return errBrokerConflict + } + if prompt.Closing || !prompt.LeaseExpiresAt.After(time.Now()) { + return errBrokerClosed + } + output = foundryResponse{ID: invocation.ResponseAlias, Status: "completed", Output: []foundryOutputItem{}} + if summary.Text != "" { + output.Output = append(output.Output, foundryOutputItem{ID: "fi_" + uuid.NewString(), Type: "message", + Content: []foundryOutputContent{{Type: "output_text", Text: summary.Text}}}) + } + link := session.Responses[invocation.ResponseAlias] + link.CallIDs = map[string]string{} + seen := map[string]bool{} + for _, item := range summary.FunctionCalls { + if !acpSafeString(item.CallID, maxProviderIdentifierBytes) || seen[item.CallID] { + return errBrokerConflict + } + seen[item.CallID] = true + alias := "fc_" + uuid.NewString() + link.CallIDs[alias] = item.CallID + item.CallID, item.ID = alias, "fi_"+uuid.NewString() + output.Output = append(output.Output, item) + } + link.Completed, link.HasFunctions = true, len(link.CallIDs) != 0 + session.Responses[invocation.ResponseAlias] = link + invocation.State, prompt.LastAlias = "completed", invocation.ResponseAlias + return nil + }) + return output, err +} + +func (b *lifecycleBroker) finishInvocation(c brokerContext, invocationError error) { + key := brokerJSONDigest(c.Owner) + b.mu.Lock() + _ = b.commitLocked(func(next *brokerLedger) error { + prompt := next.Sessions[key].Prompts[c.promptKey()] + invocation := prompt.Invocations[c.InvocationSequence] + if invocationError != nil { + prompt.Closing = true + if invocation.State == "reserved" { + invocation.State = "rejected" + } + if invocation.State == "intent" { + if errors.Is(invocationError, errBrokerAmbiguous) { + invocation.State = "uncertain" + } else { + invocation.State = "rejected" + } + } + } + return nil + }) + delete(b.active, key) + b.mu.Unlock() + if invocationError != nil { + b.startReconcile(key, true) + } +} + +func (b *lifecycleBroker) closePrompt(key, promptKey string) { + b.mu.Lock() + if b.ctx.Err() == nil { + _ = b.commitLocked(func(next *brokerLedger) error { + if session := next.Sessions[key]; session != nil { + if prompt := session.Prompts[promptKey]; prompt != nil { + prompt.Closing = true + } + } + return nil + }) + if active, ok := b.active[key]; ok && active.prompt == promptKey { + active.cancel() + } + } + b.mu.Unlock() + b.startReconcile(key, true) +} diff --git a/broker_retirement_capacity_test.go b/broker_retirement_capacity_test.go new file mode 100644 index 0000000..6d5da18 --- /dev/null +++ b/broker_retirement_capacity_test.go @@ -0,0 +1,263 @@ +package main + +import ( + "context" + "encoding/json" + "fmt" + "net/http" + "path/filepath" + "testing" + "time" +) + +const brokerHistoricalOperationCapacity = 16384 + +func brokerFillHistoricalOperationCapacity(t *testing.T, b *lifecycleBroker, c brokerContext, count int) (brokerContext, map[string]string) { + t.Helper() + renewal, body := brokerTestControlContext(brokerRenewPath, c) + renewal.BodySHA256 = brokerSHA(body) + b.mu.Lock() + err := b.commitLocked(func(next *brokerLedger) error { + session := next.Sessions[brokerJSONDigest(c.Owner)] + prompt := session.Prompts[c.promptKey()] + for len(session.Operations) < count { + renewal.OperationID = fmt.Sprintf("historical-renewal-%d", len(session.Operations)) + renewal.LeaseGeneration = prompt.LeaseGeneration + 1 + session.Operations[renewal.OperationID] = brokerOperationDigest(brokerRenewPath, renewal) + prompt.LeaseGeneration = renewal.LeaseGeneration + } + return nil + }) + valid := brokerLedgerValid(b.ledger, b.cfg.configDigest) + operations := make(map[string]string, len(b.ledger.Sessions[brokerJSONDigest(c.Owner)].Operations)) + for key, value := range b.ledger.Sessions[brokerJSONDigest(c.Owner)].Operations { + operations[key] = value + } + b.mu.Unlock() + if err != nil || !valid { + t.Fatal("saturated historical fixture did not preserve valid durable ownership") + } + return renewal, operations +} + +func TestBrokerOperationCapacityPreservesSettlementAndRetirement(t *testing.T) { + for _, settleFirst := range []bool{false, true} { + name := "direct-retirement" + if settleFirst { + name = "settlement-then-retirement" + } + t.Run(name, func(t *testing.T) { + f := newBrokerFixture(t, "success") + cfg := brokerTestConfig(t, f) + b, server := startBrokerTest(t, cfg) + c := brokerTestContext(cfg) + c.LeaseExpiresAt = time.Now().Add(2 * time.Minute).UTC().Format(time.RFC3339Nano) + status, _, err := brokerTestHTTP(context.Background(), server.URL, brokerResponsesPath, c, brokerTestBody("")) + if err != nil || status != http.StatusOK { + t.Fatal("initial fixture inference failed") + } + renewal, operations := brokerFillHistoricalOperationCapacity(t, b, c, brokerHistoricalOperationCapacity) + fresh := renewal + fresh.OperationID = "renewal-after-capacity" + fresh.LeaseGeneration++ + status, _, err = brokerTestHTTP(context.Background(), server.URL, brokerRenewPath, fresh, []byte("{}")) + if err != nil || status != http.StatusServiceUnavailable { + t.Fatal("saturated owner admitted more ordinary operation records") + } + status, _, err = brokerTestHTTP(context.Background(), server.URL, brokerRenewPath, renewal, []byte("{}")) + if err != nil || status != http.StatusOK { + t.Fatal("saturation rejected an exact recorded renewal duplicate") + } + status, _, err = brokerTestHTTP(context.Background(), server.URL, brokerResponsesPath, c, brokerTestBody("")) + if err != nil || status != http.StatusConflict { + t.Fatal("saturation replayed a recorded inference") + } + conflict, body := brokerTestControlContext(brokerRetirePath, c) + conflict.OperationID = renewal.OperationID + status, _, err = brokerTestHTTP(context.Background(), server.URL, brokerRetirePath, conflict, body) + if err != nil || status != http.StatusConflict { + t.Fatal("cleanup capacity bypassed a recorded operation conflict") + } + if settleFirst { + proof := brokerTestControl(t, server.URL, brokerSettlePath, c) + if !proof.SettlementProven { + t.Fatal("saturated owner could not settle its current prompt") + } + extra, body := brokerTestControlContext(brokerSettlePath, c) + extra.OperationID = "extra-settlement-after-capacity" + status, _, err := brokerTestHTTP(context.Background(), server.URL, brokerSettlePath, extra, body) + if err != nil || status != http.StatusServiceUnavailable { + t.Fatal("extra settlement consumed retirement capacity") + } + } + proof := brokerTestControl(t, server.URL, brokerRetirePath, c) + if !proof.RetirementProven || proof.OwnerDigest != brokerJSONDigest(c.Owner) || !brokerDigestValid(proof.ProofDigest) { + t.Fatal("saturated owner lacked exact durable retirement proof") + } + // Distinct cleanup IDs remain bounded; the original request stays + // replayable even after the two reserved slots are occupied. + for i := range 3 { + extra, body := brokerTestControlContext(brokerRetirePath, c) + extra.OperationID = fmt.Sprintf("extra-retirement-%d", i) + status, _, err := brokerTestHTTP(context.Background(), server.URL, brokerRetirePath, extra, body) + expected := http.StatusServiceUnavailable + if !settleFirst && i == 0 { + expected = http.StatusOK + } + if err != nil || status != expected { + t.Fatal("retirement operation capacity was not bounded") + } + } + b.mu.Lock() + session := b.ledger.Sessions[brokerJSONDigest(c.Owner)] + bounded := len(session.Operations) == brokerHistoricalOperationCapacity+2 + retained := true + for key, value := range operations { + retained = retained && session.Operations[key] == value + } + valid := brokerLedgerValid(b.ledger, cfg.configDigest) + b.mu.Unlock() + if !bounded || !retained || !valid { + t.Fatal("cleanup capacity lost idempotency records or invalidated the bounded ledger") + } + b.close() + server.Close() + _, server = startBrokerTest(t, cfg) + reopened := brokerTestControl(t, server.URL, brokerRetirePath, c) + if !reopened.RetirementProven || reopened.ProofDigest != proof.ProofDigest { + t.Fatal("saturated retired owner did not reopen with the same proof") + } + creates, inferences, stops, deletes := f.counts() + if creates != 1 || inferences != 1 || stops != 0 || deletes != 1 { + t.Fatal("capacity recovery replayed work or repeated remote deletion") + } + }) + } +} + +func TestBrokerOperationCapacityRetainsUnknownCreation(t *testing.T) { + f := newBrokerFixture(t, "late-create") + cfg := brokerTestConfig(t, f) + b, server := startBrokerTest(t, cfg) + c := brokerTestContext(cfg) + c.LeaseExpiresAt = time.Now().Add(2 * time.Minute).UTC().Format(time.RFC3339Nano) + _ = brokerTestControl(t, server.URL, brokerRenewPath, c) + renewed, _ := brokerFillHistoricalOperationCapacity(t, b, c, brokerHistoricalOperationCapacity-1) + c.LeaseGeneration = renewed.LeaseGeneration + c.LeaseExpiresAt = renewed.LeaseExpiresAt + done := brokerAsyncInference(context.Background(), server.URL, c, brokerTestBody("")) + select { + case <-f.started: + case <-time.After(3 * time.Second): + t.Fatal("last ordinary operation did not attempt the original creation") + } + _ = brokerWaitInference(t, done) + retirement, body := brokerTestControlContext(brokerRetirePath, c) + status, data, err := brokerTestHTTP(context.Background(), server.URL, brokerRetirePath, retirement, body) + var proof brokerControlResponse + if err != nil || status != http.StatusConflict || json.Unmarshal(data, &proof) != nil || + !proof.CreatePending || proof.RetirementProven || proof.SettlementProven || proof.ProofDigest != "" { + t.Fatalf("saturated unknown creation did not retain pending retirement: status=%d", status) + } + creates, inferences, _, deletes := f.counts() + if creates != 1 || inferences != 0 || deletes != 0 { + t.Fatal("capacity handling replayed or deleted an unknown creation") + } + f.unblock() + brokerAwait(t, func() bool { + f.mu.Lock() + defer f.mu.Unlock() + return len(f.sessions) == 1 + }) + // Reserved cleanup capacity cannot manufacture the original CREATE ack. + brokerPendingControl(t, server.URL, brokerRetirePath, c, true, 0) + creates, inferences, _, deletes = f.counts() + if creates != 1 || inferences != 0 || deletes != 0 { + t.Fatal("saturated unacknowledged creation was replayed or retired") + } +} + +func TestBrokerOperationCapacityCancelsAcknowledgedInvocation(t *testing.T) { + f := newBrokerFixture(t, "hold-known") + cfg := brokerTestConfig(t, f) + b, server := startBrokerTest(t, cfg) + c := brokerTestContext(cfg) + c.LeaseExpiresAt = time.Now().Add(2 * time.Minute).UTC().Format(time.RFC3339Nano) + done := brokerAsyncInference(context.Background(), server.URL, c, brokerTestBody("")) + brokerAwait(t, func() bool { return brokerInvocationState(b, c) == "accepted" }) + _, _ = brokerFillHistoricalOperationCapacity(t, b, c, brokerHistoricalOperationCapacity) + settlement, body := brokerTestControlContext(brokerSettlePath, c) + status, _, err := brokerTestHTTP(context.Background(), server.URL, brokerSettlePath, settlement, body) + if err != nil || (status != http.StatusOK && status != http.StatusConflict) { + t.Fatal("saturated owner rejected cancellation of an active invocation") + } + result := brokerWaitInference(t, done) + if result.err == nil && result.status == http.StatusOK { + t.Fatal("cancelled active invocation exposed a terminal result") + } + var firstProof string + for i := range 3 { + proof := brokerTestControl(t, server.URL, brokerSettlePath, c) + if !proof.SettlementProven || proof.ActiveInvocations != 0 || proof.AmbiguousInvocations != 0 { + t.Fatal("repeated cancellation lost its settlement proof") + } + if i == 0 { + firstProof = proof.ProofDigest + } else if proof.ProofDigest != firstProof { + t.Fatal("repeated cancellation replaced its durable settlement proof") + } + extra := settlement + extra.OperationID = fmt.Sprintf("extra-cancellation-%d", i) + status, _, err := brokerTestHTTP(context.Background(), server.URL, brokerSettlePath, extra, body) + if err != nil || status != http.StatusServiceUnavailable { + t.Fatal("distinct repeated cancellation consumed retirement capacity") + } + } + proof := brokerTestControl(t, server.URL, brokerRetirePath, c) + creates, inferences, stops, deletes := f.counts() + if !proof.RetirementProven || creates != 1 || inferences != 1 || stops != 1 || deletes != 1 { + t.Fatal("saturated acknowledged invocation lost containment or retirement ownership") + } + b.mu.Lock() + session := b.ledger.Sessions[brokerJSONDigest(c.Owner)] + valid := len(session.Operations) == brokerHistoricalOperationCapacity+2 && brokerLedgerValid(b.ledger, cfg.configDigest) + b.mu.Unlock() + if !valid { + t.Fatal("repeated cancellation grew or invalidated the durable ledger") + } +} + +func TestBrokerOperationCapacityRejectsOversizedLedger(t *testing.T) { + cfg := brokerConfiguration{configDigest: brokerSHA([]byte("bounded-capacity-fixture")), stateDir: filepath.Join(t.TempDir(), "broker")} + c := brokerTestContext(cfg) + store, ledger, err := openBrokerStore(cfg.stateDir, cfg.configDigest) + if err != nil { + t.Fatal("could not create bounded ledger fixture") + } + session := &brokerSession{Owner: c.Owner, CreateState: "none", Prompts: map[string]*brokerPrompt{}, + Responses: map[string]brokerResponseID{}, Operations: map[string]string{}} + ledger.Sessions[brokerJSONDigest(c.Owner)] = session + for i := range brokerHistoricalOperationCapacity + 2 { + session.Operations[fmt.Sprintf("historical-operation-%d", i)] = brokerSHA([]byte(fmt.Sprintf("operation-%d", i))) + } + err = store.save(ledger) + store.close() + if err != nil { + t.Fatal("could not save maximum ledger fixture") + } + store, ledger, err = openBrokerStore(cfg.stateDir, cfg.configDigest) + if err != nil { + t.Fatal("recovery rejected the maximum operation count") + } + ledger.Sessions[brokerJSONDigest(c.Owner)].Operations["over-capacity"] = brokerSHA([]byte("extra-operation")) + err = store.save(ledger) + store.close() + if err != nil { + t.Fatal("could not save oversized ledger fixture") + } + store, _, err = openBrokerStore(cfg.stateDir, cfg.configDigest) + if err == nil { + store.close() + t.Fatal("recovery accepted an oversized operation ledger") + } +} diff --git a/broker_review_lifecycle_test.go b/broker_review_lifecycle_test.go new file mode 100644 index 0000000..b34259a --- /dev/null +++ b/broker_review_lifecycle_test.go @@ -0,0 +1,281 @@ +package main + +import ( + "bytes" + "context" + "encoding/json" + "net/http" + "os" + "path/filepath" + "sync" + "testing" + "time" +) + +func brokerReviewFaultStore(t *testing.T, dir string) (string, func()) { + t.Helper() + retained := dir + "-retained" + if err := os.Rename(dir, retained); err != nil { + t.Fatal("could not retain fixture ledger before storage fault") + } + restored := false + restore := func() { + if restored { + return + } + if err := os.Remove(dir); err != nil && !os.IsNotExist(err) { + t.Fatal("could not remove fixture storage fault") + } + if err := os.Rename(retained, dir); err != nil { + t.Fatal("could not restore retained fixture ledger") + } + restored = true + } + t.Cleanup(restore) + // A regular file at the directory path makes every attempted save fail + // deterministically, including when the tests run with root privileges. + if err := os.WriteFile(dir, []byte("fixture storage unavailable"), 0o600); err != nil { + t.Fatal("could not install fixture storage fault") + } + return retained, restore +} + +func TestBrokerStorageFailureCancelsActiveRequests(t *testing.T) { + for _, trigger := range []string{"expiry", "settle", "other-owner-renewal"} { + t.Run(trigger, func(t *testing.T) { + f := newBrokerFixture(t, "hold-known") + cfg := brokerTestConfig(t, f) + b, server := startBrokerTest(t, cfg) + c := brokerTestContext(cfg) + if trigger == "expiry" { + c.LeaseExpiresAt = time.Now().Add(800 * time.Millisecond).UTC().Format(time.RFC3339Nano) + } + ctx, cancel := context.WithCancel(context.Background()) + defer cancel() + done := brokerAsyncInference(ctx, server.URL, c, brokerTestBody("")) + brokerAwait(t, func() bool { return brokerInvocationState(b, c) == "accepted" }) + before, err := os.ReadFile(filepath.Join(cfg.stateDir, "state.json")) + if err != nil { + t.Fatal("could not read durable fixture ownership") + } + retained, restore := brokerReviewFaultStore(t, cfg.stateDir) + if trigger != "expiry" { + path, closing := brokerSettlePath, c + if trigger == "other-owner-renewal" { + path = brokerRenewPath + closing.Owner.RuntimeSessionUID = "fixture-other-session" + } + closing, body := brokerTestControlContext(path, closing) + status, _, err := brokerTestHTTP(context.Background(), server.URL, path, closing, body) + if err != nil || status != http.StatusServiceUnavailable { + t.Fatal("control did not report the storage failure") + } + } + brokerAwait(t, func() bool { + response, err := http.Get(server.URL + "/healthz") + if err != nil { + return false + } + _ = response.Body.Close() + return response.StatusCode == http.StatusServiceUnavailable + }) + select { + case result := <-done: + if result.err == nil && result.status == http.StatusOK { + t.Fatal("storage failure exposed a successful response") + } + case <-time.After(time.Second): + t.Fatal("storage failure left accepted inference running") + } + after, err := os.ReadFile(filepath.Join(retained, "state.json")) + if err != nil || !bytes.Equal(before, after) { + t.Fatal("storage fault changed the last durable ownership record") + } + b.mu.Lock() + active, poisoned := len(b.active), b.storageError != nil + b.mu.Unlock() + creates, inferences, stops, deletes := f.counts() + if active != 0 || !poisoned || creates != 1 || inferences != 1 || stops != 0 || deletes != 0 { + t.Fatal("poisoned storage kept active authority or claimed remote cleanup") + } + closing, body := brokerTestControlContext(brokerSettlePath, c) + status, _, err := brokerTestHTTP(context.Background(), server.URL, brokerSettlePath, closing, body) + if err != nil || status != http.StatusServiceUnavailable { + t.Fatal("poisoned storage exposed settlement proof") + } + b.close() + server.Close() + restore() + _, restarted := startBrokerTest(t, cfg) + proof := brokerTestControl(t, restarted.URL, brokerSettlePath, c) + if !proof.SettlementProven || proof.AmbiguousInvocations != 0 { + t.Fatal("recovery lost the original acknowledged owner") + } + proof = brokerTestControl(t, restarted.URL, brokerRetirePath, c) + creates, inferences, stops, deletes = f.counts() + if !proof.RetirementProven || creates != 1 || inferences != 1 || stops == 0 || deletes != 1 { + t.Fatal("storage recovery replayed work or skipped acknowledged cleanup") + } + }) + } +} + +func TestBrokerRetirementRejectsNewSettlementPrompt(t *testing.T) { + for _, state := range []string{"retiring", "retired"} { + t.Run(state, func(t *testing.T) { + f := newBrokerFixture(t, "success") + cfg := brokerTestConfig(t, f) + deleteStarted, releaseDelete := make(chan struct{}), make(chan struct{}) + var deleteOnce, releaseOnce sync.Once + release := func() { releaseOnce.Do(func() { close(releaseDelete) }) } + defer release() + client := &http.Client{Transport: brokerFixtureTransport(func(r *http.Request) (*http.Response, error) { + if state == "retiring" && r.Method == http.MethodDelete { + deleteOnce.Do(func() { close(deleteStarted) }) + select { + case <-releaseDelete: + case <-r.Context().Done(): + return nil, r.Context().Err() + } + } + return http.DefaultTransport.RoundTrip(r) + })} + b, server := startBrokerTestWithClient(t, cfg, client) + c := brokerTestContext(cfg) + status, _, err := brokerTestHTTP(context.Background(), server.URL, brokerResponsesPath, c, brokerTestBody("")) + if err != nil || status != http.StatusOK { + t.Fatal("initial fixture inference failed") + } + _ = brokerTestControl(t, server.URL, brokerSettlePath, c) + if state == "retiring" { + retire, body := brokerTestControlContext(brokerRetirePath, c) + status, _, err := brokerTestHTTP(context.Background(), server.URL, brokerRetirePath, retire, body) + if err != nil || status != http.StatusConflict { + t.Fatal("retirement did not wait for deletion acknowledgement") + } + select { + case <-deleteStarted: + case <-time.After(time.Second): + t.Fatal("retirement did not reach the held deletion") + } + } else { + _ = brokerTestControl(t, server.URL, brokerRetirePath, c) + } + before, err := os.ReadFile(filepath.Join(cfg.stateDir, "state.json")) + if err != nil { + t.Fatal("could not read fixture retirement ownership") + } + later := c + later.TaskUID, later.PromptID = "fixture-task-after-retirement", "fixture-prompt-after-retirement" + later, body := brokerTestControlContext(brokerSettlePath, later) + status, _, err = brokerTestHTTP(context.Background(), server.URL, brokerSettlePath, later, body) + if err != nil || status != http.StatusGone { + t.Errorf("new settlement identity was admitted after retirement began: status=%d", status) + } + after, err := os.ReadFile(filepath.Join(cfg.stateDir, "state.json")) + if err != nil || !bytes.Equal(before, after) { + t.Error("rejected settlement identity changed durable retirement ownership") + } + b.mu.Lock() + valid := brokerLedgerValid(b.ledger, cfg.configDigest) + prompts := len(b.ledger.Sessions[brokerJSONDigest(c.Owner)].Prompts) + b.mu.Unlock() + if !valid || prompts != 1 { + t.Errorf("retirement accepted new prompt or invalidated ledger: valid=%v prompts=%d", valid, prompts) + } + proof := brokerTestControl(t, server.URL, brokerSettlePath, c) + if !proof.SettlementProven { + t.Fatal("existing settlement lost idempotent proof during retirement") + } + release() + _ = brokerTestControl(t, server.URL, brokerRetirePath, c) + b.close() + server.Close() + _, restarted := startBrokerTest(t, cfg) + proof = brokerTestControl(t, restarted.URL, brokerRetirePath, c) + if !proof.RetirementProven { + t.Fatal("retired ledger did not reopen with the original proof") + } + }) + } +} + +func TestBrokerCancelledResponseRetainsAcceptanceEvidence(t *testing.T) { + for _, shape := range []string{"json", "sse-terminal-only", "sse-after-created"} { + t.Run(shape, func(t *testing.T) { + media := "application/json" + if shape != "json" { + media = "text/event-stream" + } + f := newBrokerEvidenceFixture(t, func(request foundryResponseRequest) (string, []byte) { + response := map[string]any{"id": "provider-cancelled", "status": "cancelled", "agent_session_id": request.AgentSessionID, "output": []any{}} + var value any = response + if media == "text/event-stream" { + value = map[string]any{"type": "response.cancelled", "response": response} + } + data, _ := json.Marshal(value) + if media == "text/event-stream" { + data = []byte("data: " + string(data) + "\n\n") + } + if shape == "sse-after-created" { + response["status"] = "in_progress" + created, _ := json.Marshal(map[string]any{"type": "response.created", "response": response}) + data = append([]byte("data: "+string(created)+"\n\n"), data...) + } + return media, data + }) + cfg := brokerTestConfig(t, f) + b, server := startBrokerTest(t, cfg) + c := brokerTestContext(cfg) + status, _, err := brokerTestHTTP(context.Background(), server.URL, brokerResponsesPath, c, brokerTestBody("")) + if err != nil || status == http.StatusOK { + t.Fatal("cancelled provider response was exposed as success") + } + if state := brokerInvocationState(b, c); state != "accepted" && state != "settled" { + t.Fatalf("coherent cancellation lost acceptance evidence: state=%s", state) + } + proof := brokerTestControl(t, server.URL, brokerSettlePath, c) + if !proof.SettlementProven || proof.AmbiguousInvocations != 0 { + t.Fatal("acknowledged cancellation could not prove stop containment") + } + proof = brokerTestControl(t, server.URL, brokerRetirePath, c) + creates, inferences, stops, deletes := f.counts() + if !proof.RetirementProven || creates != 1 || inferences != 1 || stops == 0 || deletes != 1 { + t.Fatal("acknowledged cancellation replayed work or skipped cleanup") + } + }) + } +} + +func TestBrokerCancelledSSERequiresCoherentEvidence(t *testing.T) { + for _, failure := range []string{"status-mismatch", "response-error", "wrong-session"} { + t.Run(failure, func(t *testing.T) { + f := newBrokerEvidenceFixture(t, func(request foundryResponseRequest) (string, []byte) { + response := map[string]any{"id": "provider-cancelled", "status": "cancelled", "agent_session_id": request.AgentSessionID, "output": []any{}} + switch failure { + case "status-mismatch": + response["status"] = "in_progress" + case "response-error": + response["error"] = map[string]string{"code": "server_error", "message": "fixture-error-do-not-persist"} + case "wrong-session": + response["agent_session_id"] = "fixture-other-session" + } + data, _ := json.Marshal(map[string]any{"type": "response.cancelled", "response": response}) + return "text/event-stream", []byte("data: " + string(data) + "\n\n") + }) + cfg := brokerTestConfig(t, f) + b, server := startBrokerTest(t, cfg) + c := brokerTestContext(cfg) + status, _, err := brokerTestHTTP(context.Background(), server.URL, brokerResponsesPath, c, brokerTestBody("")) + if err != nil || status == http.StatusOK || brokerInvocationState(b, c) != "uncertain" { + t.Fatal("incoherent cancellation acknowledged an invocation") + } + brokerPendingControl(t, server.URL, brokerSettlePath, c, false, 1) + brokerPendingControl(t, server.URL, brokerRetirePath, c, false, 1) + _, _, _, deletes := f.counts() + if deletes != 0 { + t.Fatal("incoherent cancellation authorized deletion") + } + }) + } +} diff --git a/broker_schema_case_test.go b/broker_schema_case_test.go new file mode 100644 index 0000000..69ffa46 --- /dev/null +++ b/broker_schema_case_test.go @@ -0,0 +1,184 @@ +package main + +import ( + "bytes" + "context" + "encoding/base64" + "encoding/json" + "net/http" + "net/http/httptest" + "os" + "path/filepath" + "sync/atomic" + "testing" +) + +const brokerCaseToolSchema = `[{"type":"function","name":"hosted-probe-read","parameters":{"type":"object","properties":{},"additionalProperties":false}}]` + +func brokerCaseRequest(t *testing.T, b *lifecycleBroker, c brokerContext, members string) *httptest.ResponseRecorder { + t.Helper() + body := brokerTestBody("") + if members != "" { + body = append([]byte("{"+members+","), body[1:]...) + } + c.BodySHA256 = brokerSHA(body) + raw, err := json.Marshal(c) + if err != nil { + t.Fatal("could not encode fixture context") + } + request := httptest.NewRequest(http.MethodPost, brokerResponsesPath, bytes.NewReader(body)) + request.Header.Set("Authorization", "Bearer "+brokerFixtureBearer) + request.Header.Set(brokerContextHeader, base64.RawURLEncoding.EncodeToString(raw)) + response := httptest.NewRecorder() + b.ServeHTTP(response, request) + return response +} + +func brokerCaseRejectBeforeOwnership(t *testing.T, mode, members string) { + t.Helper() + f := newBrokerFixture(t, "success") + cfg := brokerTestConfig(t, f) + cfg.agent.ToolSchemaMode = mode + cfg.configDigest = brokerJSONDigest(cfg.agent) + var tokenCalls, transportCalls atomic.Int64 + provider := brokerEvidenceTokenProvider(func(context.Context) (string, error) { + tokenCalls.Add(1) + return brokerTestToken(), nil + }) + client := &http.Client{Transport: brokerFixtureTransport(func(r *http.Request) (*http.Response, error) { + transportCalls.Add(1) + return http.DefaultTransport.RoundTrip(r) + })} + b, err := newLifecycleBroker(context.Background(), cfg, provider, client) + if err != nil { + t.Fatal("could not start protected-field fixture") + } + t.Cleanup(b.close) + path := filepath.Join(cfg.stateDir, "state.json") + before, err := os.ReadFile(path) + if err != nil { + t.Fatal("could not capture initial fixture ledger") + } + response := brokerCaseRequest(t, b, brokerTestContext(cfg), members) + if response.Code != http.StatusBadRequest { + t.Errorf("protected field was not rejected: status=%d", response.Code) + } + if tokenCalls.Load() != 0 || transportCalls.Load() != 0 { + t.Errorf("protected field reached authentication or transport: tokenCalls=%d transportCalls=%d", tokenCalls.Load(), transportCalls.Load()) + } + b.mu.Lock() + owners, active := len(b.ledger.Sessions), len(b.active) + b.mu.Unlock() + if owners != 0 || active != 0 { + t.Errorf("protected field reserved ownership: owners=%d active=%d", owners, active) + } + after, err := os.ReadFile(path) + if err != nil || !bytes.Equal(before, after) { + t.Error("protected field changed the durable ledger") + } + creates, inferences, stops, deletes := f.counts() + if creates+inferences+stops+deletes != 0 { + t.Error("protected field reached a provider mutation") + } +} + +func TestBrokerProtectedFieldsRejectDecoderCaseVariantsBeforeOwnership(t *testing.T) { + fields := []struct { + name string + keys []string + values []string + modes []string + }{ + {"tools", []string{`"tools"`, `"Tools"`, `"TOOLS"`, `"tOoLs"`, `"toolſ"`, `"TOOLſ"`, `"\u0054ools"`, `"tool\u017f"`}, + []string{brokerCaseToolSchema, `[]`, `null`}, []string{toolSchemaModeProviderStatic}}, + {"agent_session_id", []string{`"agent_session_id"`, `"Agent_Session_ID"`, `"AGENT_SESSION_ID"`, `"aGeNt_sEsSiOn_iD"`, `"agent_ſeſſion_id"`, `"\u0041gent_session_id"`, `"agent_\u017fe\u017f\u017fion_id"`}, + []string{`"fixture-injected-session"`, `""`, `null`}, []string{toolSchemaModeProviderStatic, toolSchemaModeRequest}}, + } + for _, field := range fields { + for _, mode := range field.modes { + for _, key := range field.keys { + for i, value := range field.values { + t.Run(field.name+"/"+mode+"/"+key+"/"+[]string{"value", "empty", "null"}[i], func(t *testing.T) { + brokerCaseRejectBeforeOwnership(t, mode, key+":"+value) + }) + } + } + } + } +} + +func TestBrokerProtectedFieldsRejectDuplicatePresenceBeforeOwnership(t *testing.T) { + for _, tc := range []struct { + name string + members string + modes []string + }{ + {"tools_folded_null_last", `"Tools":` + brokerCaseToolSchema + `,"TOOLS":null`, []string{toolSchemaModeProviderStatic}}, + {"tools_folded_value_last", `"Tools":null,"TOOLS":` + brokerCaseToolSchema, []string{toolSchemaModeProviderStatic}}, + {"tools_folded_empty_last", `"Tools":` + brokerCaseToolSchema + `,"TOOLſ":[]`, []string{toolSchemaModeProviderStatic}}, + {"tools_canonical_null", `"tools":null,"Tools":` + brokerCaseToolSchema, []string{toolSchemaModeProviderStatic}}, + {"tools_exact_duplicate", `"Tools":null,"Tools":[]`, []string{toolSchemaModeProviderStatic}}, + {"tools_escaped_duplicate", `"Tools":null,"\u0054ools":[]`, []string{toolSchemaModeProviderStatic}}, + {"session_folded_null_last", `"Agent_Session_ID":"fixture-injected-session","AGENT_SESSION_ID":null`, []string{toolSchemaModeProviderStatic, toolSchemaModeRequest}}, + {"session_folded_value_last", `"Agent_Session_ID":null,"AGENT_SESSION_ID":"fixture-injected-session"`, []string{toolSchemaModeProviderStatic, toolSchemaModeRequest}}, + {"session_folded_empty_last", `"Agent_Session_ID":"fixture-injected-session","agent_ſeſſion_id":""`, []string{toolSchemaModeProviderStatic, toolSchemaModeRequest}}, + {"session_canonical_null", `"agent_session_id":null,"Agent_Session_ID":""`, []string{toolSchemaModeProviderStatic, toolSchemaModeRequest}}, + {"session_exact_duplicate", `"Agent_Session_ID":null,"Agent_Session_ID":""`, []string{toolSchemaModeProviderStatic, toolSchemaModeRequest}}, + {"session_escaped_duplicate", `"Agent_Session_ID":null,"\u0041gent_Session_ID":""`, []string{toolSchemaModeProviderStatic, toolSchemaModeRequest}}, + } { + for _, mode := range tc.modes { + t.Run(tc.name+"/"+mode, func(t *testing.T) { + brokerCaseRejectBeforeOwnership(t, mode, tc.members) + }) + } + } +} + +func TestBrokerRequestToolCaseVariantsRemainValid(t *testing.T) { + for _, tc := range []struct { + name string + mode string + members string + tools int + }{ + {"static_omitted", toolSchemaModeProviderStatic, "", 0}, + {"request_omitted", toolSchemaModeRequest, "", 0}, + {"canonical", toolSchemaModeRequest, `"tools":` + brokerCaseToolSchema, 1}, + {"title", toolSchemaModeRequest, `"Tools":` + brokerCaseToolSchema, 1}, + {"upper", toolSchemaModeRequest, `"TOOLS":` + brokerCaseToolSchema, 1}, + {"mixed", toolSchemaModeRequest, `"tOoLs":` + brokerCaseToolSchema, 1}, + {"unicode_fold", toolSchemaModeRequest, `"toolſ":` + brokerCaseToolSchema, 1}, + {"escaped", toolSchemaModeRequest, `"\u0054ools":` + brokerCaseToolSchema, 1}, + {"escaped_unicode_fold", toolSchemaModeRequest, `"tool\u017f":` + brokerCaseToolSchema, 1}, + {"empty", toolSchemaModeRequest, `"Tools":[]`, 0}, + {"null", toolSchemaModeRequest, `"TOOLS":null`, 0}, + } { + t.Run(tc.name, func(t *testing.T) { + f := newBrokerFixture(t, "success") + cfg := brokerTestConfig(t, f) + cfg.agent.ToolSchemaMode = tc.mode + cfg.configDigest = brokerJSONDigest(cfg.agent) + b, server := startBrokerTest(t, cfg) + c := brokerTestContext(cfg) + if response := brokerCaseRequest(t, b, c, tc.members); response.Code != http.StatusOK { + t.Fatalf("valid tool mode request failed: status=%d", response.Code) + } + f.mu.Lock() + if len(f.requests) != 1 || len(f.requests[0].Tools) != tc.tools { + t.Error("valid request did not preserve its tool schemas") + } else if tc.tools != 0 { + encoded, err := json.Marshal(f.requests[0].Tools) + if err != nil || string(encoded) != brokerCaseToolSchema { + t.Error("valid request changed its tool schema") + } + } + f.mu.Unlock() + creates, inferences, _, _ := f.counts() + if creates != 1 || inferences != 1 { + t.Error("valid request did not submit exactly one owned inference") + } + _ = brokerTestControl(t, server.URL, brokerSettlePath, c) + _ = brokerTestControl(t, server.URL, brokerRetirePath, c) + }) + } +} diff --git a/broker_settlement_test.go b/broker_settlement_test.go new file mode 100644 index 0000000..dfcdb7b --- /dev/null +++ b/broker_settlement_test.go @@ -0,0 +1,85 @@ +package main + +import ( + "bytes" + "context" + "encoding/json" + "io" + "net/http" + "net/http/httptest" + "strings" + "sync/atomic" + "testing" +) + +func TestBrokerCompletedSettlementPreservesRemoteConversation(t *testing.T) { + for _, restart := range []bool{false, true} { + name := "same-broker" + if restart { + name = "recovered-broker" + } + t.Run(name, func(t *testing.T) { + f := newBrokerFixture(t, "success") + f.server.Close() + var historyLost atomic.Bool + // Foundry stop terminates compute. Model a backend that retains + // response history in RAM and rejects continuation after a stop. + f.server = httptest.NewServer(http.HandlerFunc(func(w http.ResponseWriter, r *http.Request) { + if r.Method == http.MethodPost && strings.HasSuffix(r.URL.Path, ":stop") { + historyLost.Store(true) + } + if r.Method == http.MethodPost && strings.HasSuffix(r.URL.Path, "/protocols/openai/responses") { + raw, err := io.ReadAll(r.Body) + var request foundryResponseRequest + if err != nil || json.Unmarshal(raw, &request) != nil { + w.WriteHeader(http.StatusBadRequest) + return + } + r.Body = io.NopCloser(bytes.NewReader(raw)) + if request.PreviousResponseID != "" && historyLost.Load() { + w.WriteHeader(http.StatusBadRequest) + return + } + } + f.serve(w, r) + })) + cfg := brokerTestConfig(t, f) + b, server := startBrokerTest(t, cfg) + c := brokerTestContext(cfg) + status, data, err := brokerTestHTTP(context.Background(), server.URL, brokerResponsesPath, c, brokerTestBody("")) + var first foundryResponse + if err != nil || status != http.StatusOK || json.Unmarshal(data, &first) != nil { + t.Fatalf("seed response failed: status=%d", status) + } + if restart { + b.close() + server.Close() + b, server = startBrokerTest(t, cfg) + } + proof := brokerTestControl(t, server.URL, brokerSettlePath, c) + if !proof.SettlementProven || proof.ActiveInvocations != 0 || !brokerDigestValid(proof.ProofDigest) { + t.Fatal("completed request lacked durable settlement") + } + c.TaskUID, c.PromptID, c.OperationID = "continuation-task", "continuation-prompt", "continuation-request" + status, _, err = brokerTestHTTP(context.Background(), server.URL, brokerResponsesPath, c, brokerTestBody(first.ID)) + if err != nil || status != http.StatusOK { + t.Fatalf("completed settlement destroyed remote conversation: status=%d", status) + } + _ = brokerTestControl(t, server.URL, brokerSettlePath, c) + proof = brokerTestControl(t, server.URL, brokerRetirePath, c) + if !proof.RetirementProven { + t.Fatal("completed session retirement lacked deletion proof") + } + creates, inferences, stops, deletes := f.counts() + if creates != 1 || inferences != 2 || stops != 0 || deletes != 1 { + t.Fatalf("unexpected lifecycle operations: %d %d %d %d", creates, inferences, stops, deletes) + } + b.mu.Lock() + valid := brokerLedgerValid(b.ledger, cfg.configDigest) + b.mu.Unlock() + if !valid { + t.Fatal("completed settlement invalidated ownership evidence") + } + }) + } +} diff --git a/broker_storage_poison_test.go b/broker_storage_poison_test.go new file mode 100644 index 0000000..58c46ff --- /dev/null +++ b/broker_storage_poison_test.go @@ -0,0 +1,102 @@ +package main + +import ( + "bytes" + "context" + "errors" + "net/http" + "os" + "path/filepath" + "strings" + "sync" + "testing" + "time" +) + +func TestBrokerStoragePoisonCancelsDetachedMutations(t *testing.T) { + for _, phase := range []string{"create", "stop", "delete"} { + t.Run(phase, func(t *testing.T) { + mode := "success" + if phase == "stop" { + mode = "hold-known" + } + f := newBrokerFixture(t, mode) + cfg := brokerTestConfig(t, f) + cfg.operationTimeout = 5 * time.Second + entered := make(chan context.Context, 1) + release := make(chan struct{}) + unblock := sync.OnceFunc(func() { close(release) }) + defer unblock() + client := &http.Client{Transport: brokerFixtureTransport(func(r *http.Request) (*http.Response, error) { + selected := phase == "create" && r.Method == http.MethodPost && strings.HasSuffix(r.URL.Path, "/endpoint/sessions") || + phase == "stop" && r.Method == http.MethodPost && strings.HasSuffix(r.URL.Path, ":stop") || + phase == "delete" && r.Method == http.MethodDelete + if selected { + entered <- r.Context() + select { + case <-r.Context().Done(): + return nil, r.Context().Err() + case <-release: + } + } + return http.DefaultTransport.RoundTrip(r) + })} + b, server := startBrokerTestWithClient(t, cfg, client) + c := brokerTestContext(cfg) + done := brokerAsyncInference(context.Background(), server.URL, c, brokerTestBody("")) + if phase == "stop" { + brokerAwait(t, func() bool { return brokerInvocationState(b, c) == "accepted" }) + b.closePrompt(brokerJSONDigest(c.Owner), c.promptKey()) + _ = brokerWaitInference(t, done) + } else if phase == "delete" { + result := brokerWaitInference(t, done) + if result.err != nil || result.status != http.StatusOK { + t.Fatal("original invocation did not complete") + } + brokerPendingControl(t, server.URL, brokerRetirePath, c, false, 0) + } + var mutationCtx context.Context + select { + case mutationCtx = <-entered: + case <-time.After(3 * time.Second): + t.Fatal("selected original mutation did not enter transport") + } + retained, restore := brokerReviewFaultStore(t, cfg.stateDir) + before, err := os.ReadFile(filepath.Join(retained, "state.json")) + if err != nil { + t.Fatal("could not read original intent") + } + other := c + other.Owner.RuntimeSessionUID = "other-owner" + other, body := brokerTestControlContext(brokerRenewPath, other) + status, _, err := brokerTestHTTP(context.Background(), server.URL, brokerRenewPath, other, body) + if err != nil || status != http.StatusServiceUnavailable { + t.Fatal("storage poison was not observed") + } + cancelled := mutationCtx.Err() != nil + unblock() + if phase == "create" { + result := brokerWaitInference(t, done) + if result.err == nil && result.status == http.StatusOK { + t.Fatal("poisoned broker exposed successful output") + } + } + brokerAwait(t, func() bool { + b.mu.Lock() + defer b.mu.Unlock() + return len(b.workers) == 0 + }) + b.close() + server.Close() + after, err := os.ReadFile(filepath.Join(retained, "state.json")) + if err != nil || !bytes.Equal(before, after) || !errors.Is(b.storageError, errBrokerStorage) { + t.Fatal("storage poison changed durable ownership or disappeared") + } + restore() + creates, inferences, stops, deletes := f.counts() + if !cancelled || stops != 0 || deletes != 0 || phase == "create" && creates+inferences != 0 { + t.Fatalf("detached mutation continued after poison: cancelled=%v creates=%d inferences=%d stops=%d deletes=%d", cancelled, creates, inferences, stops, deletes) + } + }) + } +} diff --git a/broker_store.go b/broker_store.go new file mode 100644 index 0000000..2fb40b4 --- /dev/null +++ b/broker_store.go @@ -0,0 +1,185 @@ +package main + +import ( + "encoding/json" + "errors" + "io" + "os" + "path/filepath" + "strings" + "syscall" + "time" + + "github.com/google/uuid" +) + +const brokerMaxLedgerBytes = 32 << 20 + +// The supervisor freezes one settlement operation per prompt and one +// retirement operation per owner, retrying those exact contexts. Before each +// new invocation or ordinary growth, retain space for every nonretired owner: +// two maximally escaped 4096-byte response IDs, two maximally escaped 512-byte +// cleanup keys with their digests, and UUID, alias, link, proof and state growth. +// Only one prompt (and one live invocation) can be unsettled for an owner. +// Completion call maps are not reserved; they must fit before output is exposed. +// Unbounded distinct cleanup IDs from a lifecycle bearer holder are outside this +// caller guarantee; their operation history is still retained and hard-capped. +const brokerOwnerReserveBytes = 64 << 10 +const brokerPrincipalReserveBytes = 128 + +var errBrokerCapacity = errors.New("Foundry broker durable state is at capacity") + +func brokerLedgerReserveBytes(ledger *brokerLedger) int { + reserve := 0 + if ledger.PrincipalDigest == "" { + reserve = brokerPrincipalReserveBytes + } + for _, session := range ledger.Sessions { + if !session.Retired { + reserve += brokerOwnerReserveBytes + } + } + return reserve +} + +type brokerLedger struct { + Version uint32 `json:"version"` + ConfigDigest string `json:"configDigest"` + PrincipalDigest string `json:"principalDigest,omitempty"` + Sessions map[string]*brokerSession `json:"sessions"` +} + +type brokerSession struct { + Owner brokerOwner `json:"owner"` + RemoteID string `json:"remoteID,omitempty"` + CreateState string `json:"createState"` + Retiring bool `json:"retiring"` + Retired bool `json:"retired"` + ProofDigest string `json:"proofDigest,omitempty"` + CurrentPrompt string `json:"currentPrompt,omitempty"` + Prompts map[string]*brokerPrompt `json:"prompts"` + Responses map[string]brokerResponseID `json:"responses"` + Operations map[string]string `json:"operations"` +} + +type brokerPrompt struct { + Identity brokerContext `json:"identity"` + LeaseGeneration uint64 `json:"leaseGeneration"` + LeaseExpiresAt time.Time `json:"leaseExpiresAt"` + Closing bool `json:"closing"` + Settled bool `json:"settled"` + ProofDigest string `json:"proofDigest,omitempty"` + LastSequence uint64 `json:"lastSequence"` + LastAlias string `json:"lastAlias,omitempty"` + Invocations map[uint64]*brokerInvocation `json:"invocations"` +} + +type brokerInvocation struct { + Sequence uint64 `json:"sequence"` + OperationID string `json:"operationID"` + BodyDigest string `json:"bodyDigest"` + State string `json:"state"` + ResponseID string `json:"responseID,omitempty"` + ResponseAlias string `json:"responseAlias,omitempty"` +} + +type brokerResponseID struct { + RemoteID string `json:"remoteID"` + PromptKey string `json:"promptKey"` + Completed bool `json:"completed"` + HasFunctions bool `json:"hasFunctions"` + CallIDs map[string]string `json:"callIDs,omitempty"` +} + +type brokerStore struct { + dir string + lock *os.File +} + +func openBrokerStore(dir, digest string) (*brokerStore, *brokerLedger, error) { + if !filepath.IsAbs(dir) || filepath.Clean(dir) == string(filepath.Separator) || !brokerDigestValid(digest) { + return nil, nil, errBrokerStorage + } + lock, created, err := openStoreLock(dir, "broker.lock", syncStoreDirectory) + if err != nil { + return nil, nil, errBrokerStorage + } + store := &brokerStore{dir: dir, lock: lock} + ledger := &brokerLedger{Version: 1, ConfigDigest: digest, Sessions: map[string]*brokerSession{}} + path := filepath.Join(dir, "state.json") + info, err := os.Lstat(path) + if os.IsNotExist(err) { + if !created { + store.close() + return nil, nil, errBrokerStorage + } + if err := store.save(ledger); err != nil { + store.close() + return nil, nil, err + } + return store, ledger, nil + } + if err != nil || !brokerPrivateFile(info) || info.Size() > brokerMaxLedgerBytes { + store.close() + return nil, nil, errBrokerStorage + } + data, err := os.ReadFile(path) + if err != nil || acpDecode(data, ledger, true) != nil || !brokerLedgerValid(ledger, digest) { + store.close() + return nil, nil, errBrokerStorage + } + return store, ledger, nil +} + +func brokerPrivateFile(info os.FileInfo) bool { + if info == nil || !info.Mode().IsRegular() || info.Mode().Perm() != 0o600 { + return false + } + stat, ok := info.Sys().(*syscall.Stat_t) + return ok && int(stat.Uid) == os.Geteuid() && stat.Nlink == 1 +} + +func (s *brokerStore) save(ledger *brokerLedger) error { + data, err := json.Marshal(ledger) + if err != nil { + return errBrokerStorage + } + return s.saveBytes(data) +} + +func (s *brokerStore) saveBytes(data []byte) error { + if len(data) > brokerMaxLedgerBytes { + return errBrokerCapacity + } + name := filepath.Join(s.dir, ".state-"+uuid.NewString()) + file, err := os.OpenFile(name, os.O_WRONLY|os.O_CREATE|os.O_EXCL|syscall.O_NOFOLLOW, 0o600) + if err != nil { + return errBrokerStorage + } + defer func() { _ = os.Remove(name) }() + _, err = io.Copy(file, strings.NewReader(string(data))) + if err == nil { + err = file.Sync() + } + closeErr := file.Close() + if err != nil || closeErr != nil || os.Rename(name, filepath.Join(s.dir, "state.json")) != nil { + return errBrokerStorage + } + directory, err := os.Open(s.dir) + if err != nil { + return errBrokerStorage + } + err = directory.Sync() + closeErr = directory.Close() + if err != nil || closeErr != nil { + return errBrokerStorage + } + return nil +} + +func (s *brokerStore) close() { + if s != nil && s.lock != nil { + _ = syscall.Flock(int(s.lock.Fd()), syscall.LOCK_UN) + _ = s.lock.Close() + } +} diff --git a/broker_store_test.go b/broker_store_test.go new file mode 100644 index 0000000..442e50a --- /dev/null +++ b/broker_store_test.go @@ -0,0 +1,204 @@ +package main + +import ( + "bytes" + "context" + "encoding/json" + "net/http" + "net/http/httptest" + "os" + "path/filepath" + "strings" + "sync/atomic" + "testing" +) + +func TestBrokerStoreRejectsCorruptOwnership(t *testing.T) { + f := newBrokerFixture(t, "functions") + cfg := brokerTestConfig(t, f) + b, server := startBrokerTest(t, cfg) + c := brokerTestContext(cfg) + status, _, err := brokerTestHTTP(context.Background(), server.URL, brokerResponsesPath, c, brokerTestBody("")) + if err != nil || status != http.StatusOK { + t.Fatal("fixture ownership was not created") + } + b.close() + server.Close() + baseline, err := os.ReadFile(filepath.Join(cfg.stateDir, "state.json")) + if err != nil { + t.Fatal("fixture ledger unavailable") + } + for _, kind := range []string{"unknown_state", "missing_response", "orphan_response", "response_completion", "owner_fence", + "missing_operation", "operation_digest", "last_sequence", "last_alias", "premature_proof", "premature_retirement", + "missing_principal", "missing_current_prompt", "prompt_identity", "duplicate_member"} { + t.Run(kind, func(t *testing.T) { + var ledger brokerLedger + if json.Unmarshal(baseline, &ledger) != nil { + t.Fatal("invalid baseline fixture") + } + session := ledger.Sessions[brokerJSONDigest(c.Owner)] + prompt := session.Prompts[c.promptKey()] + invocation := prompt.Invocations[1] + switch kind { + case "unknown_state": + invocation.State = "looks-finished" + case "missing_response": + delete(session.Responses, invocation.ResponseAlias) + case "orphan_response": + session.Responses["orphan"] = session.Responses[invocation.ResponseAlias] + case "response_completion": + link := session.Responses[invocation.ResponseAlias] + link.Completed = false + session.Responses[invocation.ResponseAlias] = link + case "owner_fence": + session.Owner.ControllerEpoch++ + case "missing_operation": + delete(session.Operations, invocation.OperationID) + case "operation_digest": + session.Operations[invocation.OperationID] = "broken" + case "last_sequence": + prompt.LastSequence++ + case "last_alias": + prompt.LastAlias = "" + case "premature_proof": + prompt.ProofDigest = brokerSHA([]byte("invented")) + case "premature_retirement": + session.Retired = true + case "missing_principal": + ledger.PrincipalDigest = "" + case "missing_current_prompt": + session.CurrentPrompt = "" + case "prompt_identity": + prompt.Identity.PromptRequestDigest = "broken" + } + data, _ := json.Marshal(ledger) + if kind == "duplicate_member" { + data = append([]byte(`{"version":1,`), data[1:]...) + } + dir := filepath.Join(t.TempDir(), "broker") + if os.Mkdir(dir, 0o700) != nil || os.WriteFile(filepath.Join(dir, "state.json"), data, 0o600) != nil { + t.Fatal("could not prepare corrupt ledger fixture") + } + store, _, err := openBrokerStore(dir, cfg.configDigest) + if err == nil { + store.close() + t.Fatal("corrupt state was accepted for recovery") + } + }) + } + store, _, err := openBrokerStore(cfg.stateDir, cfg.configDigest) + if err != nil { + t.Fatal("unchanged valid ownership was rejected") + } + store.close() +} + +func TestBrokerStorePrivatePermissionsAndSingleWriter(t *testing.T) { + digest := brokerSHA([]byte("private-store-test")) + dir := filepath.Join(t.TempDir(), "broker") + store, _, err := openBrokerStore(dir, digest) + if err != nil { + t.Fatal("private ledger initialization failed") + } + other, _, err := openBrokerStore(dir, digest) + if err == nil { + other.close() + t.Fatal("two durable writers acquired ownership") + } + store.close() + if os.Chmod(filepath.Join(dir, "state.json"), 0o644) != nil { + t.Fatal("permission fixture failed") + } + store, _, err = openBrokerStore(dir, digest) + if err == nil { + store.close() + t.Fatal("publicly readable ledger was accepted") + } + if os.Chmod(filepath.Join(dir, "state.json"), 0o600) != nil || os.Chmod(dir, 0o755) != nil { + t.Fatal("directory permission fixture failed") + } + store, _, err = openBrokerStore(dir, digest) + if err == nil { + store.close() + t.Fatal("nonprivate state directory was accepted") + } +} + +func TestBrokerPersistenceFailureClosesAdmissionAndHealth(t *testing.T) { + f := newBrokerFixture(t, "success") + cfg := brokerTestConfig(t, f) + b, server := startBrokerTest(t, cfg) + if os.Rename(cfg.stateDir, cfg.stateDir+"-moved") != nil { + t.Fatal("could not inject persistence failure") + } + c := brokerTestContext(cfg) + for range 2 { + status, _, err := brokerTestHTTP(context.Background(), server.URL, brokerResponsesPath, c, brokerTestBody("")) + if err != nil || status != http.StatusServiceUnavailable { + t.Fatal("persistence failure did not close inference admission") + } + } + request := httptest.NewRequest(http.MethodGet, "/healthz", nil) + response := httptest.NewRecorder() + b.ServeHTTP(response, request) + if response.Code != http.StatusServiceUnavailable { + t.Fatal("poisoned ledger was reported healthy") + } + creates, inferences, _, _ := f.counts() + if creates != 0 || inferences != 0 { + t.Fatal("remote work preceded durable ownership") + } +} + +func TestBrokerHealthCLIRequiresNoConfigurationCredentialOrLock(t *testing.T) { + for _, status := range []int{http.StatusOK, http.StatusServiceUnavailable} { + t.Run(http.StatusText(status), func(t *testing.T) { + var calls atomic.Int32 + server := httptest.NewServer(http.HandlerFunc(func(w http.ResponseWriter, r *http.Request) { + calls.Add(1) + if r.Method != http.MethodGet || r.URL.Path != "/healthz" || r.Header.Get("Authorization") != "" { + t.Error("health checker expanded authority") + } + w.WriteHeader(status) + })) + defer server.Close() + t.Setenv("ORKA_FOUNDRY_BROKER_ADDR", strings.TrimPrefix(server.URL, "http://")) + t.Setenv("ORKA_FOUNDRY_BROKER_STATE_DIR", "/not-used-for-health") + handled, err := maybeServeBroker([]string{"--protocol", "broker", "--health-check", "--config", "/missing.json"}) + if !handled || calls.Load() != 1 || (err == nil) != (status == http.StatusOK) { + t.Fatal("health command initialized config/auth or misclassified readiness") + } + }) + } + for _, address := range []string{"example.com:80", "10.0.0.1:8091", "0.0.0.0:8091", "127.0.0.1:0", "127.0.0.1:65536"} { + if checkBrokerHealth(address) == nil { + t.Fatal("health checker accepted a nonlocal or invalid target") + } + } +} + +func TestBrokerEmptyRetirementIsDurableAdmissionTombstone(t *testing.T) { + f := newBrokerFixture(t, "success") + cfg := brokerTestConfig(t, f) + b, server := startBrokerTest(t, cfg) + c := brokerTestContext(cfg) + proof := brokerTestControl(t, server.URL, brokerRetirePath, c) + if !proof.RetirementProven || proof.RemoteSessionCreated { + t.Fatal("no-inference retirement lacked proof") + } + b.close() + server.Close() + _, server = startBrokerTest(t, cfg) + status, _, err := brokerTestHTTP(context.Background(), server.URL, brokerResponsesPath, c, brokerTestBody("")) + if err != nil || status != http.StatusGone { + t.Fatal("restarted tombstone admitted delayed inference") + } + data, err := os.ReadFile(filepath.Join(cfg.stateDir, "state.json")) + if err != nil || bytes.Contains(data, []byte("fixture-input-do-not-persist")) { + t.Fatal("retirement persisted child input") + } + creates, inferences, stops, deletes := f.counts() + if creates+inferences+stops+deletes != 0 { + t.Fatal("empty tombstone caused a remote operation") + } +} diff --git a/broker_store_validation.go b/broker_store_validation.go new file mode 100644 index 0000000..a3bc5d7 --- /dev/null +++ b/broker_store_validation.go @@ -0,0 +1,154 @@ +package main + +import ( + "strings" + "time" + + "github.com/google/uuid" +) + +// Recovery never interprets an unknown state or a broken response reference as +// quiescence. Content is deliberately absent; these are ownership records only. +func brokerLedgerValid(ledger *brokerLedger, digest string) bool { + if ledger.Version != 1 || ledger.ConfigDigest != digest || ledger.Sessions == nil || len(ledger.Sessions) > 4096 || + (ledger.PrincipalDigest != "" && !brokerDigestValid(ledger.PrincipalDigest)) { + return false + } + for key, session := range ledger.Sessions { + if !brokerSessionValid(session, key, digest) || (session.RemoteID != "" && ledger.PrincipalDigest == "") { + return false + } + } + return true +} + +func brokerSessionValid(session *brokerSession, key, digest string) bool { + if session == nil || !session.Owner.valid() || key != brokerJSONDigest(session.Owner) || + session.Prompts == nil || session.Responses == nil || session.Operations == nil || + len(session.Prompts) > 4096 || len(session.Operations) > brokerRetirementOperationLimit { + return false + } + switch session.CreateState { + case "none": + if session.RemoteID != "" || len(session.Responses) != 0 { + return false + } + case "intent", "known", "deleted": + id, err := uuid.Parse(session.RemoteID) + if err != nil || id.String() != session.RemoteID { + return false + } + default: + return false + } + if session.Retired { + if !session.Retiring || (session.CreateState != "none" && session.CreateState != "deleted") || !brokerDigestValid(session.ProofDigest) { + return false + } + } else if session.CreateState == "deleted" || session.ProofDigest != "" { + return false + } + if (len(session.Prompts) == 0) != (session.CurrentPrompt == "") || + (session.CurrentPrompt != "" && session.Prompts[session.CurrentPrompt] == nil) { + return false + } + for id, operation := range session.Operations { + if !acpSafeString(id, 512) || !brokerDigestValid(operation) { + return false + } + } + linked := make(map[string]bool, len(session.Responses)) + for promptKey, prompt := range session.Prompts { + if !brokerPromptValid(prompt, promptKey, digest, session, linked) { + return false + } + } + return len(linked) == len(session.Responses) +} + +func brokerPromptValid(prompt *brokerPrompt, key, digest string, session *brokerSession, linked map[string]bool) bool { + if prompt == nil || prompt.Identity.promptKey() != key || prompt.Identity.Owner != session.Owner || + prompt.Identity.Protocol != brokerProtocol || prompt.Identity.AgentConfigurationDigest != digest || + !acpSafeString(prompt.Identity.TaskUID, 512) || prompt.Identity.TaskAttempt == 0 || + !acpSafeString(prompt.Identity.PromptID, 512) || !brokerDigestValid(prompt.Identity.PromptRequestDigest) || + !brokerDigestValid(prompt.Identity.BodySHA256) || session.Operations[prompt.Identity.OperationID] == "" || + prompt.Invocations == nil || prompt.Identity.LeaseGeneration == 0 || prompt.LeaseGeneration < prompt.Identity.LeaseGeneration { + return false + } + expiry, err := time.Parse(time.RFC3339Nano, prompt.Identity.LeaseExpiresAt) + if err != nil || prompt.LeaseExpiresAt.Before(expiry) || prompt.LeaseExpiresAt.IsZero() { + return false + } + if prompt.Settled { + if !prompt.Closing || !brokerDigestValid(prompt.ProofDigest) { + return false + } + } else if prompt.ProofDigest != "" || session.Retired || key != session.CurrentPrompt { + return false + } + if session.Retiring && !prompt.Closing { + return false + } + var last, lastCompleted uint64 + lastAlias := "" + for sequence, invocation := range prompt.Invocations { + if !brokerInvocationValid(invocation, sequence, key, prompt, session, linked) { + return false + } + if sequence > last { + last = sequence + } + if invocation.State == "completed" && sequence > lastCompleted { + lastCompleted, lastAlias = sequence, invocation.ResponseAlias + } + } + return prompt.LastSequence == last && prompt.LastAlias == lastAlias +} + +func brokerInvocationValid(invocation *brokerInvocation, sequence uint64, promptKey string, prompt *brokerPrompt, session *brokerSession, linked map[string]bool) bool { + if invocation == nil || invocation.Sequence != sequence || sequence == 0 || !brokerDigestValid(invocation.BodyDigest) || + !acpSafeString(invocation.OperationID, 512) || session.Operations[invocation.OperationID] == "" { + return false + } + switch invocation.State { + case "reserved", "intent", "rejected", "uncertain": + if invocation.ResponseID != "" || invocation.ResponseAlias != "" || prompt.Settled { + return false + } + case "accepted", "completed": + if invocation.ResponseID == "" || (prompt.Settled && invocation.State != "completed") { + return false + } + case "settled": + if !prompt.Settled { + return false + } + default: + return false + } + if invocation.ResponseID == "" { + return invocation.ResponseAlias == "" + } + link, ok := session.Responses[invocation.ResponseAlias] + if !ok || linked[invocation.ResponseAlias] || !brokerAliasValid(invocation.ResponseAlias, "fr_") || + validateProviderIdentifier("response", link.RemoteID) != nil || link.RemoteID != invocation.ResponseID || + link.PromptKey != promptKey || link.Completed != (invocation.State == "completed") || + link.HasFunctions != (len(link.CallIDs) > 0) || (!link.Completed && len(link.CallIDs) != 0) { + return false + } + for alias, remote := range link.CallIDs { + if !brokerAliasValid(alias, "fc_") || validateProviderIdentifier("call", remote) != nil { + return false + } + } + linked[invocation.ResponseAlias] = true + return true +} + +func brokerAliasValid(value, prefix string) bool { + if !strings.HasPrefix(value, prefix) { + return false + } + id, err := uuid.Parse(strings.TrimPrefix(value, prefix)) + return err == nil && prefix+id.String() == value +} diff --git a/broker_test.go b/broker_test.go new file mode 100644 index 0000000..7db004c --- /dev/null +++ b/broker_test.go @@ -0,0 +1,546 @@ +package main + +import ( + "bytes" + "context" + "encoding/base64" + "encoding/json" + "fmt" + "io" + "net/http" + "net/http/httptest" + "os" + "path/filepath" + "strings" + "sync" + "testing" + "time" +) + +const brokerFixtureBearer = "test-broker-bearer-with-more-than-thirty-two-bytes" + +type brokerFixtureToken struct{ value string } + +func (p brokerFixtureToken) AccessToken(context.Context) (string, error) { return p.value, nil } + +type brokerFixtureTransport func(*http.Request) (*http.Response, error) + +func (f brokerFixtureTransport) RoundTrip(r *http.Request) (*http.Response, error) { return f(r) } + +func brokerTestToken() string { + return "fixture." + base64.RawURLEncoding.EncodeToString([]byte(`{"aud":"https://ai.azure.com","tid":"test-tenant","oid":"test-principal","appid":"test-client"}`)) + ".fixture" +} + +type brokerFixture struct { + t *testing.T + server *httptest.Server + mu sync.Mutex + sessions map[string]string + creates int + inferences int + stops int + deletes int + mode string + started chan struct{} + release chan struct{} + startOnce sync.Once + releaseOnce sync.Once + requests []foundryResponseRequest + createCheck func(string) +} + +func newBrokerFixture(t *testing.T, mode string) *brokerFixture { + t.Helper() + f := &brokerFixture{t: t, sessions: map[string]string{}, mode: mode, started: make(chan struct{}), release: make(chan struct{})} + f.server = httptest.NewServer(http.HandlerFunc(f.serve)) + t.Cleanup(func() { f.unblock(); f.server.Close() }) + return f +} + +func (f *brokerFixture) unblock() { f.releaseOnce.Do(func() { close(f.release) }) } + +func (f *brokerFixture) counts() (int, int, int, int) { + f.mu.Lock() + defer f.mu.Unlock() + return f.creates, f.inferences, f.stops, f.deletes +} + +func (f *brokerFixture) serve(w http.ResponseWriter, r *http.Request) { + if r.Header.Get("Authorization") != "Bearer "+brokerTestToken() { + f.t.Error("provider authentication mismatch") + w.WriteHeader(401) + return + } + if r.URL.Query().Get("api-version") != "v1" { + f.t.Error("provider API version mismatch") + w.WriteHeader(400) + return + } + prefix := "/api/projects/fixture/agents/fixture" + w.Header().Set("Content-Type", "application/json") + suffix := strings.TrimPrefix(r.URL.Path, prefix) + if suffix == "" && r.Method == http.MethodGet { + _, _ = io.WriteString(w, `{"name":"fixture","agent_endpoint":{"authorization_schemes":[{"type":"entra"}]}}`) + return + } + if suffix == "/versions/3" && r.Method == http.MethodGet { + _, _ = io.WriteString(w, `{"name":"fixture","version":"3","status":"active","definition":{"kind":"hosted"}}`) + return + } + if suffix == "/endpoint/sessions" && r.Method == http.MethodPost { + var request brokerRemoteSession + if json.NewDecoder(r.Body).Decode(&request) != nil || request.ID == "" || request.Version.Type != "version_ref" || request.Version.Version != "3" { + f.t.Error("session creation did not contain exact chosen identity and version") + w.WriteHeader(400) + return + } + f.mu.Lock() + f.creates++ + check := f.createCheck + f.mu.Unlock() + if check != nil { + check(request.ID) + } + if strings.HasPrefix(f.mode, "create-rejected") { + status := http.StatusTooManyRequests + switch f.mode { + case "create-rejected-server": + status = http.StatusServiceUnavailable + case "create-rejected-conflict": + status = http.StatusConflict + case "create-rejected-truncated": + w.Header().Set("Content-Length", "1024") + } + w.WriteHeader(status) + if f.mode == "create-rejected-oversized" { + _, _ = io.WriteString(w, strings.Repeat("x", acpMaxConfigBytes+1)) + return + } + _, _ = io.WriteString(w, `{"error":"fixture creation rejection"}`) + return + } + if f.mode == "late-create" { + // Lose the original acknowledgement independently of prompt + // cancellation, then let that same request create the session later. + conn, _, err := w.(http.Hijacker).Hijack() + if err != nil { + f.t.Error("could not drop creation acknowledgement") + return + } + _ = conn.Close() + f.startOnce.Do(func() { close(f.started) }) + <-f.release + } + f.mu.Lock() + if _, exists := f.sessions[request.ID]; exists { + f.mu.Unlock() + w.WriteHeader(409) + return + } + f.sessions[request.ID] = "active" + f.mu.Unlock() + if f.mode == "late-create" { + return + } + if f.mode == "hold-create" { + f.startOnce.Do(func() { close(f.started) }) + <-f.release + } + w.WriteHeader(201) + f.sessionJSON(w, request.ID, "active") + return + } + if suffix == "/endpoint/protocols/openai/responses" && r.Method == http.MethodPost { + var request foundryResponseRequest + if json.NewDecoder(r.Body).Decode(&request) != nil { + w.WriteHeader(400) + return + } + f.mu.Lock() + _, exists := f.sessions[request.AgentSessionID] + if !exists { + f.mu.Unlock() + w.WriteHeader(404) + return + } + f.sessions[request.AgentSessionID] = "active" + f.inferences++ + count := f.inferences + f.requests = append(f.requests, request) + f.mu.Unlock() + responseID := fmt.Sprintf("provider-response-%d", count) + if f.mode == "hold-known" || f.mode == "hold-unknown" || f.mode == "truncated" { + w.Header().Set("Content-Type", "text/event-stream") + if f.mode != "hold-unknown" { + frame, _ := json.Marshal(map[string]any{"type": "response.created", "response": map[string]any{ + "id": responseID, "status": "in_progress", "agent_session_id": request.AgentSessionID, "output": []any{}}}) + _, _ = fmt.Fprintf(w, "data: %s\n\n", frame) + } + w.(http.Flusher).Flush() + f.startOnce.Do(func() { close(f.started) }) + if f.mode == "truncated" { + return + } + select { + case <-r.Context().Done(): + return + case <-f.release: + } + frame, _ := json.Marshal(map[string]any{"type": "response.completed", "response": map[string]any{ + "id": responseID, "status": "completed", "agent_session_id": request.AgentSessionID, + "output": []any{map[string]any{"type": "message", "role": "assistant", "content": []any{map[string]any{"type": "output_text", "text": "fixture-terminal-do-not-persist"}}}}}}) + _, _ = fmt.Fprintf(w, "data: %s\n\n", frame) + return + } + if strings.HasPrefix(f.mode, "rejected") { + status := http.StatusTooManyRequests + if f.mode == "rejected-server" { + status = http.StatusServiceUnavailable + } + if f.mode == "rejected-truncated" { + w.Header().Set("Content-Length", "1024") + } + w.WriteHeader(status) + if f.mode == "rejected-oversized" { + _, _ = io.WriteString(w, strings.Repeat("x", acpMaxConfigBytes+1)) + return + } + _, _ = io.WriteString(w, `{"error":"fixture rejection"}`) + return + } + if f.mode == "functions" && count == 1 { + _ = json.NewEncoder(w).Encode(map[string]any{"id": responseID, "status": "completed", "agent_session_id": request.AgentSessionID, + "output": []any{map[string]any{"id": "provider-item-id", "type": "function_call", "call_id": "provider-call-id", "name": "fixture_echo", "arguments": "{\"text\":\"fixture\"}"}}}) + return + } + _ = json.NewEncoder(w).Encode(map[string]any{"id": responseID, "status": "completed", "agent_session_id": request.AgentSessionID, + "output": []any{map[string]any{"id": "provider-item-id", "type": "message", "role": "assistant", "content": []any{map[string]any{"type": "output_text", "text": "fixture-terminal-do-not-persist"}}}}}) + return + } + if strings.HasPrefix(suffix, "/endpoint/sessions/") { + id := strings.TrimPrefix(suffix, "/endpoint/sessions/") + stop := strings.HasSuffix(id, ":stop") + id = strings.TrimSuffix(id, ":stop") + f.mu.Lock() + state, exists := f.sessions[id] + if stop && r.Method == http.MethodPost { + f.stops++ + if !exists { + f.mu.Unlock() + w.WriteHeader(404) + return + } + f.sessions[id] = "idle" + f.mu.Unlock() + if state == "idle" { + w.WriteHeader(409) + } else { + w.WriteHeader(204) + } + return + } + if r.Method == http.MethodDelete { + f.deletes++ + delete(f.sessions, id) + f.mu.Unlock() + w.WriteHeader(204) + return + } + f.mu.Unlock() + if r.Method == http.MethodGet { + if !exists { + w.WriteHeader(404) + return + } + f.sessionJSON(w, id, state) + return + } + } + w.WriteHeader(404) +} + +func (f *brokerFixture) sessionJSON(w http.ResponseWriter, id, state string) { + _ = json.NewEncoder(w).Encode(map[string]any{"agent_session_id": id, "version_indicator": map[string]string{"type": "version_ref", "agent_version": "3"}, "status": state}) +} + +func brokerTestConfig(t *testing.T, fixture *brokerFixture) brokerConfiguration { + t.Helper() + agent := acpAgentConfiguration{Model: "fixture-model", ToolSchemaMode: toolSchemaModeProviderStatic, + HostedTarget: acpHostedTarget{ProjectEndpoint: fixture.server.URL + "/api/projects/fixture", AgentName: "fixture", AgentVersion: "3"}} + raw, _ := json.Marshal(agent) + return brokerConfiguration{agent: agent, configDigest: brokerSHA(raw), stateDir: filepath.Join(t.TempDir(), "broker"), bearer: brokerFixtureBearer, operationTimeout: 2 * time.Second} +} + +func startBrokerTest(t *testing.T, cfg brokerConfiguration) (*lifecycleBroker, *httptest.Server) { + t.Helper() + return startBrokerTestWithClient(t, cfg, nil) +} + +func startBrokerTestWithClient(t *testing.T, cfg brokerConfiguration, client *http.Client) (*lifecycleBroker, *httptest.Server) { + t.Helper() + b, err := newLifecycleBroker(context.Background(), cfg, brokerFixtureToken{brokerTestToken()}, client) + if err != nil { + t.Fatalf("new broker: %v", err) + } + server := httptest.NewServer(b) + t.Cleanup(func() { b.close(); server.Close() }) + return b, server +} + +func brokerTestContext(cfg brokerConfiguration) brokerContext { + return brokerContext{Protocol: brokerProtocol, AgentConfigurationDigest: cfg.configDigest, + Owner: brokerOwner{RuntimeInstanceID: "fixture-instance", SupervisorBootID: "fixture-boot", ControllerEpoch: 11, + RuntimePoolUID: "fixture-pool", RuntimePoolGeneration: 2, RuntimeSessionUID: "fixture-session", RuntimeSessionGeneration: 3, + RuntimeProfileDigest: "sha256:" + strings.Repeat("1", 64), ProfileDigestSchemaVersion: 1}, + TaskUID: "fixture-task", TaskAttempt: 1, PromptID: "fixture-prompt", PromptRequestDigest: "sha256:" + strings.Repeat("2", 64), + LeaseGeneration: 1, LeaseExpiresAt: time.Now().Add(10 * time.Second).UTC().Format(time.RFC3339Nano), OperationID: "fixture-inference", InvocationSequence: 1} +} + +func brokerTestBody(previous string) []byte { + request := acpResponseRequest{Model: "fixture-model", foundryResponseRequest: foundryResponseRequest{ + Input: "fixture-input-do-not-persist", Stream: true, Store: true, PreviousResponseID: previous}} + data, _ := json.Marshal(request) + return data +} + +func brokerTestHTTP(ctx context.Context, base, path string, c brokerContext, body []byte) (int, []byte, error) { + c.BodySHA256 = brokerSHA(body) + raw, _ := json.Marshal(c) + method := http.MethodPost + if path == brokerStatusPath { + method = http.MethodGet + } + request, err := http.NewRequestWithContext(ctx, method, base+path, bytes.NewReader(body)) + if err != nil { + return 0, nil, err + } + request.Header.Set("Authorization", "Bearer "+brokerFixtureBearer) + request.Header.Set(brokerContextHeader, base64.RawURLEncoding.EncodeToString(raw)) + response, err := http.DefaultClient.Do(request) + if err != nil { + return 0, nil, err + } + defer response.Body.Close() //nolint:errcheck + data, err := io.ReadAll(response.Body) + return response.StatusCode, data, err +} + +func brokerTestControlContext(path string, c brokerContext) (brokerContext, []byte) { + c.InvocationSequence = 0 + c.OperationID = "control-" + strings.TrimPrefix(path, "/internal/v1/") + "-" + c.PromptID + if path == brokerRetirePath || path == brokerStatusPath { + c.TaskUID = "" + c.TaskAttempt = 0 + c.PromptID = "" + c.PromptRequestDigest = "" + c.LeaseGeneration = 0 + c.LeaseExpiresAt = "" + } + body := []byte("{}") + if path == brokerStatusPath { + body = nil + } + return c, body +} + +func brokerTestControl(t *testing.T, base, path string, c brokerContext) brokerControlResponse { + t.Helper() + c, body := brokerTestControlContext(path, c) + deadline := time.Now().Add(4 * time.Second) + for { + status, data, err := brokerTestHTTP(context.Background(), base, path, c, body) + var proof brokerControlResponse + if err != nil || json.Unmarshal(data, &proof) != nil { + t.Fatalf("control response unreadable, status=%d", status) + } + if status == 200 { + return proof + } + if status != 409 || time.Now().After(deadline) { + t.Fatalf("control did not settle: status=%d state=%s pending=%v ambiguous=%d", status, proof.State, proof.CreatePending, proof.AmbiguousInvocations) + } + time.Sleep(20 * time.Millisecond) + } +} + +func brokerAwait(t *testing.T, condition func() bool) { + t.Helper() + deadline := time.Now().Add(4 * time.Second) + for !condition() { + if time.Now().After(deadline) { + t.Fatal("bounded condition did not become true") + } + time.Sleep(10 * time.Millisecond) + } +} + +func TestBrokerLifecycleOwnershipContinuationAndRetirement(t *testing.T) { + f := newBrokerFixture(t, "success") + cfg := brokerTestConfig(t, f) + b, server := startBrokerTest(t, cfg) + f.createCheck = func(id string) { + data, err := os.ReadFile(filepath.Join(cfg.stateDir, "state.json")) + if err != nil || !bytes.Contains(data, []byte(id)) || !bytes.Contains(data, []byte(`"createState":"intent"`)) { + t.Error("remote create preceded durable caller-chosen ownership") + } + } + c := brokerTestContext(cfg) + status, data, err := brokerTestHTTP(context.Background(), server.URL, brokerResponsesPath, c, brokerTestBody("")) + var first foundryResponse + if err != nil || status != 200 || json.Unmarshal(data, &first) != nil { + t.Fatalf("first response failed: %d", status) + } + if !strings.HasPrefix(first.ID, "fr_") || first.AgentSessionID != "" || bytes.Contains(data, []byte("provider-")) { + t.Fatal("provider identity escaped into ACP response") + } + proof := brokerTestControl(t, server.URL, brokerSettlePath, c) + if !proof.SettlementProven || proof.ActiveInvocations != 0 || proof.AmbiguousInvocations != 0 || !brokerDigestValid(proof.ProofDigest) { + t.Fatal("first settlement lacked exact proof") + } + c.TaskUID = "fixture-task-two" + c.PromptID = "fixture-prompt-two" + c.OperationID = "inference-two" + status, _, err = brokerTestHTTP(context.Background(), server.URL, brokerResponsesPath, c, brokerTestBody(first.ID)) + if err != nil || status != 200 { + t.Fatalf("continued response failed: %d", status) + } + _ = brokerTestControl(t, server.URL, brokerSettlePath, c) + proof = brokerTestControl(t, server.URL, brokerRetirePath, c) + if !proof.RetirementProven || !proof.SettlementProven || proof.CreatePending || proof.State != "retired" { + t.Fatal("retirement lacked proof") + } + duplicate := brokerTestControl(t, server.URL, brokerRetirePath, c) + if duplicate.ProofDigest != proof.ProofDigest { + t.Fatal("retirement proof changed on duplicate") + } + creates, inferences, stops, deletes := f.counts() + if creates != 1 || inferences != 2 || stops != 0 || deletes != 1 { + t.Fatalf("unexpected operation counts: %d %d %d %d", creates, inferences, stops, deletes) + } + f.mu.Lock() + if f.requests[0].AgentSessionID != f.requests[1].AgentSessionID || f.requests[1].PreviousResponseID != "provider-response-1" { + t.Error("continuation did not preserve remote ownership") + } + f.mu.Unlock() + ledger, err := os.ReadFile(filepath.Join(cfg.stateDir, "state.json")) + if err != nil || bytes.Contains(ledger, []byte("fixture-input-do-not-persist")) || bytes.Contains(ledger, []byte("fixture-terminal-do-not-persist")) || bytes.Contains(ledger, []byte(brokerTestToken())) { + t.Fatal("ledger retained content or credentials") + } + b.mu.Lock() + retired := b.ledger.Sessions[brokerJSONDigest(c.Owner)].Retired + b.mu.Unlock() + if !retired { + t.Fatal("retirement was not durable") + } +} + +func TestBrokerNoInferenceCleanupRejectsDelayedPOST(t *testing.T) { + f := newBrokerFixture(t, "success") + cfg := brokerTestConfig(t, f) + _, server := startBrokerTest(t, cfg) + c := brokerTestContext(cfg) + c.LeaseExpiresAt = time.Now().Add(-time.Second).UTC().Format(time.RFC3339Nano) + proof := brokerTestControl(t, server.URL, brokerSettlePath, c) + if !proof.SettlementProven || proof.RemoteSessionCreated { + t.Fatal("empty prompt did not obtain never-created proof") + } + c.LeaseExpiresAt = time.Now().Add(10 * time.Second).UTC().Format(time.RFC3339Nano) + c.OperationID = "delayed-inference" + status, _, err := brokerTestHTTP(context.Background(), server.URL, brokerResponsesPath, c, brokerTestBody("")) + if err != nil || status != 410 { + t.Fatalf("delayed inference was not fenced: %d", status) + } + proof = brokerTestControl(t, server.URL, brokerRetirePath, c) + if !proof.RetirementProven || proof.RemoteSessionCreated { + t.Fatal("never-created retirement missing") + } + creates, inferences, stops, deletes := f.counts() + if creates+inferences+stops+deletes != 0 { + t.Fatal("empty cleanup performed remote mutations") + } +} + +func TestBrokerRenewBeforeInference(t *testing.T) { + f := newBrokerFixture(t, "success") + cfg := brokerTestConfig(t, f) + _, server := startBrokerTest(t, cfg) + c := brokerTestContext(cfg) + c.LeaseGeneration = 4 + proof := brokerTestControl(t, server.URL, brokerRenewPath, c) + if proof.LeaseGeneration != 4 || proof.State != "open" { + t.Fatal("early renewal was not established") + } + status, _, err := brokerTestHTTP(context.Background(), server.URL, brokerResponsesPath, c, brokerTestBody("")) + if err != nil || status != 200 { + t.Fatalf("first inference after renewal failed: %d", status) + } + c.LeaseGeneration = 5 + c.LeaseExpiresAt = time.Now().Add(20 * time.Second).UTC().Format(time.RFC3339Nano) + c.OperationID = "renew-five" + cc := c + cc.InvocationSequence = 0 + cc.BodySHA256 = brokerSHA([]byte("{}")) + status, _, err = brokerTestHTTP(context.Background(), server.URL, brokerRenewPath, cc, []byte("{}")) + if err != nil || status != 200 { + t.Fatalf("later renewal failed: %d", status) + } + c.LeaseGeneration = 4 + c.LeaseExpiresAt = time.Now().Add(-time.Second).UTC().Format(time.RFC3339Nano) + proof = brokerTestControl(t, server.URL, brokerSettlePath, c) + if !proof.SettlementProven || proof.LeaseGeneration != 5 { + t.Fatal("cleanup with old lease failed or reopened authority") + } + _ = brokerTestControl(t, server.URL, brokerRetirePath, c) +} + +func TestBrokerUnknownInferenceNeverClaimsCleanup(t *testing.T) { + f := newBrokerFixture(t, "hold-unknown") + cfg := brokerTestConfig(t, f) + b, server := startBrokerTest(t, cfg) + c := brokerTestContext(cfg) + ctx, cancel := context.WithCancel(context.Background()) + done := make(chan struct{}) + go func() { + defer close(done) + _, _, _ = brokerTestHTTP(ctx, server.URL, brokerResponsesPath, c, brokerTestBody("")) + }() + select { + case <-f.started: + case <-time.After(3 * time.Second): + t.Fatal("inference not started") + } + cancel() + <-done + brokerAwait(t, func() bool { + b.mu.Lock() + defer b.mu.Unlock() + p := b.ledger.Sessions[brokerJSONDigest(c.Owner)].Prompts[c.promptKey()] + return p.Invocations[1].State == "uncertain" + }) + cc := c + cc.InvocationSequence = 0 + cc.OperationID = "unknown-settle" + for range 3 { + status, data, err := brokerTestHTTP(context.Background(), server.URL, brokerSettlePath, cc, []byte("{}")) + var proof brokerControlResponse + if err != nil || status != 409 || json.Unmarshal(data, &proof) != nil || proof.AmbiguousInvocations != 1 || proof.SettlementProven || proof.ProofDigest != "" { + t.Fatal("ambiguous inference produced a cleanup proof") + } + time.Sleep(30 * time.Millisecond) + } + cc.TaskUID = "" + cc.TaskAttempt = 0 + cc.PromptID = "" + cc.PromptRequestDigest = "" + cc.LeaseGeneration = 0 + cc.LeaseExpiresAt = "" + cc.OperationID = "unknown-retire" + status, _, err := brokerTestHTTP(context.Background(), server.URL, brokerRetirePath, cc, []byte("{}")) + if err != nil || status != 409 { + t.Fatal("ambiguous inference was retired") + } + _, inferences, _, deletes := f.counts() + if inferences != 1 || deletes != 0 { + t.Fatal("ambiguous inference replayed or lost its remote cleanup target") + } +} diff --git a/broker_transport_deadline_test.go b/broker_transport_deadline_test.go new file mode 100644 index 0000000..9d2ce64 --- /dev/null +++ b/broker_transport_deadline_test.go @@ -0,0 +1,350 @@ +package main + +import ( + "bufio" + "bytes" + "context" + "encoding/base64" + "encoding/json" + "fmt" + "io" + "net" + "net/http" + "net/http/httptest" + "path/filepath" + "strings" + "sync" + "testing" + "testing/synctest" + "time" +) + +// These tests use the production HTTP transport and real HTTP encoding over +// net.Pipe. No socket, DNS, identity provider, or live service is used. Fake +// time preserves the actual 30/45-second deadline relationship. +type brokerDeadlineRemote struct { + mu sync.Mutex + wg sync.WaitGroup + closed bool + createDelay time.Duration + responseDelay time.Duration + getStatus string + responseStatus int + sessions map[string]bool + creates, inferences, deletes, stops int + createAckWritten bool +} + +func (f *brokerDeadlineRemote) dial(context.Context, string, string) (net.Conn, error) { + f.mu.Lock() + if f.closed { + f.mu.Unlock() + return nil, net.ErrClosed + } + f.wg.Add(1) + f.mu.Unlock() + left, right := net.Pipe() + go func() { + defer f.wg.Done() + defer right.Close() + r, err := http.ReadRequest(bufio.NewReader(right)) + if err != nil { + return + } + body, err := io.ReadAll(r.Body) + _ = r.Body.Close() + if err != nil { + return + } + status, value, created := f.reply(r, body) + raw, _ := json.Marshal(value) + response := &http.Response{StatusCode: status, Proto: "HTTP/1.1", ProtoMajor: 1, ProtoMinor: 1, + Header: http.Header{"Content-Type": []string{"application/json"}}, Body: io.NopCloser(bytes.NewReader(raw)), + ContentLength: int64(len(raw)), Close: true} + err = response.Write(right) + if created && err == nil { + f.mu.Lock() + f.createAckWritten = true + f.mu.Unlock() + } + }() + return left, nil +} + +func (f *brokerDeadlineRemote) close() { + // Transport dial goroutines can outlive the cancelled request. Close + // admission before joining so a late dial cannot race the first Add + // against Wait after the previous handlers have all exited. + f.mu.Lock() + f.closed = true + f.mu.Unlock() + f.wg.Wait() +} + +func (f *brokerDeadlineRemote) reply(r *http.Request, body []byte) (int, any, bool) { + p := strings.TrimPrefix(r.URL.Path, "/api/projects/fixture/agents/fixture") + if p == "" && r.Method == http.MethodGet { + // Separate a renewal at +30s from the header timeout at +30.2s. + time.Sleep(200 * time.Millisecond) + return 200, map[string]any{"name": "fixture", "agent_endpoint": map[string]any{ + "authorization_schemes": []any{map[string]string{"type": "entra"}}}}, false + } + if p == "/versions/3" { + return 200, map[string]any{"name": "fixture", "version": "3", "status": "active", "definition": map[string]string{"kind": "hosted"}}, false + } + if p == "/endpoint/sessions" && r.Method == http.MethodPost { + var request brokerRemoteSession + _ = json.Unmarshal(body, &request) + f.mu.Lock() + f.creates++ + f.sessions[request.ID] = true + f.mu.Unlock() + time.Sleep(f.createDelay) + return 201, brokerDeadlineSession(request.ID, "active"), true + } + if p == "/endpoint/protocols/openai/responses" && r.Method == http.MethodPost { + var request foundryResponseRequest + _ = json.Unmarshal(body, &request) + f.mu.Lock() + f.inferences++ + f.mu.Unlock() + time.Sleep(f.responseDelay) + if f.responseStatus != 200 { + return f.responseStatus, map[string]string{"error": "synthetic admission rejection"}, false + } + return 200, map[string]any{"id": "synthetic-response", "status": "completed", "agent_session_id": request.AgentSessionID, + "output": []any{map[string]any{"id": "synthetic-item", "type": "message", "role": "assistant", + "content": []any{map[string]string{"type": "output_text", "text": "ok"}}}}}, false + } + if strings.HasPrefix(p, "/endpoint/sessions/") { + id := strings.TrimPrefix(p, "/endpoint/sessions/") + stop := strings.HasSuffix(id, ":stop") + id = strings.TrimSuffix(id, ":stop") + f.mu.Lock() + defer f.mu.Unlock() + if stop && r.Method == http.MethodPost { + f.stops++ + return 204, nil, false + } + if r.Method == http.MethodDelete { + f.deletes++ + delete(f.sessions, id) + return 204, nil, false + } + if !f.sessions[id] { + return 404, nil, false + } + status := f.getStatus + if f.stops > 0 { + status = "idle" + } + return 200, brokerDeadlineSession(id, status), false + } + return 404, nil, false +} + +func brokerDeadlineSession(id, status string) map[string]any { + return map[string]any{"agent_session_id": id, "version_indicator": map[string]string{"type": "version_ref", "agent_version": "3"}, "status": status} +} + +func brokerDeadlineRequest(b *lifecycleBroker, path string, c brokerContext, body []byte) *httptest.ResponseRecorder { + c.BodySHA256 = brokerSHA(body) + raw, _ := json.Marshal(c) + r := httptest.NewRequest(http.MethodPost, path, bytes.NewReader(body)) + r.Header.Set("Authorization", "Bearer "+brokerFixtureBearer) + r.Header.Set(brokerContextHeader, base64.RawURLEncoding.EncodeToString(raw)) + w := httptest.NewRecorder() + b.ServeHTTP(w, r) + return w +} + +func brokerDeadlineControl(t *testing.T, b *lifecycleBroker, path string, c brokerContext) brokerControlResponse { + t.Helper() + c, body := brokerTestControlContext(path, c) + for i := 0; i < 100; i++ { + w := brokerDeadlineRequest(b, path, c, body) + var proof brokerControlResponse + if json.Unmarshal(w.Body.Bytes(), &proof) != nil { + t.Fatal("synthetic control has no valid receipt") + } + if w.Code == 200 { + return proof + } + if w.Code != 409 { + t.Fatalf("unexpected control status %d", w.Code) + } + time.Sleep(10 * time.Millisecond) + } + t.Fatal("synthetic cleanup did not reach a receipt") + return brokerControlResponse{} +} + +func TestBrokerRemoteHeaderDeadline(t *testing.T) { + for _, tc := range []struct { + name string + create, response, wantElapsed time.Duration + getStatus string + responseStatus, wantInvocations int + wantSuccess, wantAmbiguity bool + wantCreatePending bool + }{ + {"create_ack_inside_operation_budget", 31 * time.Second, 0, 31200 * time.Millisecond, "active", 200, 1, true, false, false}, + {"create_operation_deadline_still_bounds_original", 46 * time.Second, 0, 45200 * time.Millisecond, "active", 200, 0, false, false, true}, + {"response_headers_timeout_is_ambiguous", 0, 46 * time.Second, 45200 * time.Millisecond, "active", 200, 1, false, true, false}, + } { + t.Run(tc.name, func(t *testing.T) { + synctest.Test(t, func(t *testing.T) { + started := time.Now() + f := &brokerDeadlineRemote{createDelay: tc.create, responseDelay: tc.response, getStatus: tc.getStatus, + responseStatus: tc.responseStatus, sessions: map[string]bool{}} + client := newBrokerHTTPClient() + transport := client.Transport.(*http.Transport) + transport.DialContext = f.dial + transport.DisableKeepAlives = true + agent := acpAgentConfiguration{Model: "fixture-model", ToolSchemaMode: toolSchemaModeProviderStatic, + HostedTarget: acpHostedTarget{ProjectEndpoint: "http://synthetic.invalid/api/projects/fixture", AgentName: "fixture", AgentVersion: "3"}} + raw, _ := json.Marshal(agent) + cfg := brokerConfiguration{agent: agent, configDigest: brokerSHA(raw), stateDir: filepath.Join(t.TempDir(), "broker"), bearer: brokerFixtureBearer, operationTimeout: 45 * time.Second} + b, err := newLifecycleBroker(context.Background(), cfg, brokerFixtureToken{brokerTestToken()}, client) + if err != nil { + t.Fatal("synthetic broker did not initialize") + } + shutdown := sync.OnceFunc(func() { + b.close() + transport.CloseIdleConnections() + f.close() + }) + defer shutdown() + c := brokerTestContext(cfg) + c.LeaseExpiresAt = started.Add(30 * time.Second).UTC().Format(time.RFC3339Nano) + done := make(chan *httptest.ResponseRecorder, 1) + go func() { done <- brokerDeadlineRequest(b, brokerResponsesPath, c, brokerTestBody("")) }() + for generation := uint64(2); generation <= 3; generation++ { + time.Sleep(time.Until(started.Add(time.Duration(generation-1) * 15 * time.Second))) + renew := c + renew.LeaseGeneration = generation + renew.LeaseExpiresAt = started.Add(time.Duration(generation+1) * 15 * time.Second).UTC().Format(time.RFC3339Nano) + renew.OperationID = fmt.Sprintf("synthetic-renew-%d", generation) + renew.InvocationSequence = 0 + w := brokerDeadlineRequest(b, brokerRenewPath, renew, []byte("{}")) + var proof brokerControlResponse + if w.Code != 200 || json.Unmarshal(w.Body.Bytes(), &proof) != nil || proof.State != "open" || proof.LeaseGeneration != generation { + t.Fatal("original invocation did not acknowledge exact renewal") + } + } + w := <-done + if (w.Code == http.StatusOK) != tc.wantSuccess { + t.Fatalf("unexpected invocation HTTP %d", w.Code) + } + elapsed := time.Since(started) + if elapsed != tc.wantElapsed { + t.Fatalf("unexpected fake-time failure/completion duration %s", elapsed) + } + if tc.wantAmbiguity { + brokerAwait(t, func() bool { + b.mu.Lock() + defer b.mu.Unlock() + return b.ledger.Sessions[brokerJSONDigest(c.Owner)].Prompts[c.promptKey()].Invocations[1].State == "uncertain" + }) + } + pending := tc.wantAmbiguity || tc.wantCreatePending + if pending { + for _, path := range []string{brokerSettlePath, brokerRetirePath} { + cc, body := brokerTestControlContext(path, c) + w := brokerDeadlineRequest(b, path, cc, body) + var proof brokerControlResponse + if w.Code != http.StatusConflict || json.Unmarshal(w.Body.Bytes(), &proof) != nil || + proof.CreatePending != tc.wantCreatePending || proof.SettlementProven || proof.RetirementProven || proof.ProofDigest != "" { + t.Fatal("unacknowledged work lost pending ownership or authorized cleanup proof") + } + } + } else { + proof := brokerDeadlineControl(t, b, brokerSettlePath, c) + if !proof.SettlementProven { + t.Fatal("exact synthetic prompt settlement missing") + } + proof = brokerDeadlineControl(t, b, brokerRetirePath, c) + if !proof.RetirementProven { + t.Fatal("known exact owner retirement missing") + } + } + // Stop broker workers and fixture admission, then join the original + // delayed HTTP write before inspecting the final evidence. + shutdown() + b.mu.Lock() + owner := b.ledger.Sessions[brokerJSONDigest(c.Owner)] + prompt := owner.Prompts[c.promptKey()] + inv := prompt.Invocations[1] + if prompt.LeaseGeneration != 3 || !prompt.LeaseExpiresAt.Equal(started.Add(60*time.Second)) || len(prompt.Invocations) != 1 { + t.Fatal("renewal or invocation identity changed") + } + if !tc.wantSuccess && (inv.ResponseID != "" || len(owner.Responses) != 0) { + t.Fatal("unexpected response acceptance evidence") + } + if owner.Retired == pending || tc.wantCreatePending && owner.CreateState != "intent" { + t.Fatal("retirement classification mismatch") + } + state, createState := inv.State, owner.CreateState + b.mu.Unlock() + f.mu.Lock() + defer f.mu.Unlock() + if f.creates != 1 || f.inferences != tc.wantInvocations || f.deletes != map[bool]int{true: 0, false: 1}[pending] { + t.Fatal("replay, incorrect submission, or incorrect deletion") + } + if tc.create > cfg.operationTimeout && f.createAckWritten { + t.Fatal("original Session POST acknowledgment was unexpectedly delivered") + } + t.Logf("elapsed=%s leaseGeneration=3 expiryFromInitial=30s creates=%d inferenceSubmissions=%d deletes=%d originalCreateAckWritten=%v state=%s createState=%s", elapsed, f.creates, f.inferences, f.deletes, f.createAckWritten, state, createState) + }) + }) + } +} + +func TestBrokerDeadlineRemoteShutdown(t *testing.T) { + synctest.Test(t, func(t *testing.T) { + f := &brokerDeadlineRemote{} + original, err := f.dial(context.Background(), "tcp", "synthetic.invalid") + if err != nil { + t.Fatal(err) + } + defer original.Close() + + // A transport dial may already be queued when its request is cancelled. + queued, release := make(chan struct{}), make(chan struct{}) + type dialResult struct { + conn net.Conn + err error + } + late := make(chan dialResult, 1) + go func() { + close(queued) + <-release + conn, err := f.dial(context.Background(), "tcp", "synthetic.invalid") + late <- dialResult{conn, err} + }() + <-queued + closed := make(chan struct{}) + go func() { + f.close() + close(closed) + }() + synctest.Wait() + select { + case <-closed: + t.Fatal("fixture shutdown did not join its admitted handler") + default: + } + + close(release) + result := <-late + if result.conn != nil { + _ = result.conn.Close() + t.Fatal("fixture admitted a queued dial after shutdown") + } + if result.err != net.ErrClosed { + t.Fatalf("late dial error = %v, want closed admission", result.err) + } + _ = original.Close() + <-closed + }) +} diff --git a/broker_unsent_dispatch_test.go b/broker_unsent_dispatch_test.go new file mode 100644 index 0000000..e0c5380 --- /dev/null +++ b/broker_unsent_dispatch_test.go @@ -0,0 +1,76 @@ +package main + +import ( + "context" + "errors" + "io" + "net/http" + "strings" + "sync/atomic" + "testing" +) + +type brokerDispatchBody struct { + io.Reader + closed bool +} + +func (b *brokerDispatchBody) Close() error { b.closed = true; return nil } + +func TestBrokerCancelledBeforeDispatchIsUnsent(t *testing.T) { + for _, phase := range []string{"create", "inference"} { + t.Run(phase, func(t *testing.T) { + f := newBrokerFixture(t, "success") + cfg := brokerTestConfig(t, f) + var attempts atomic.Int32 + client := &http.Client{Transport: brokerFixtureTransport(func(*http.Request) (*http.Response, error) { + attempts.Add(1) + return nil, context.Canceled + })} + b, _ := startBrokerTestWithClient(t, cfg, client) + ctx, cancel := context.WithCancel(context.Background()) + defer cancel() + request, err := b.prepareRemoteSessionCreate(ctx, brokerIdentityRemoteSession) + if err != nil { + t.Fatal("request preparation failed") + } + body := &brokerDispatchBody{Reader: strings.NewReader("{}")} + request.Body = body + cancel() + if phase == "create" { + created, createErr := b.remoteSessionCreate(request, brokerIdentityRemoteSession) + err = createErr + if created { + t.Fatal("uncalled creation fabricated acceptance") + } + } else { + _, err = b.sendRemoteRequest(request) + } + if !errors.Is(err, errBrokerRequestUnsent) || errors.Is(err, errBrokerAmbiguous) || attempts.Load() != 0 || !body.closed { + t.Fatalf("pre-dispatch cancellation was not proven unsent: attempts=%d unsent=%v ambiguous=%v bodyClosed=%v", attempts.Load(), errors.Is(err, errBrokerRequestUnsent), errors.Is(err, errBrokerAmbiguous), body.closed) + } + }) + } +} + +func TestBrokerDispatchErrorsRemainAmbiguous(t *testing.T) { + for _, transportError := range []error{context.Canceled, context.DeadlineExceeded, errBrokerRequestUnsent} { + t.Run(transportError.Error(), func(t *testing.T) { + var attempts atomic.Int32 + b := &lifecycleBroker{httpClient: &http.Client{Transport: brokerFixtureTransport(func(request *http.Request) (*http.Response, error) { + attempts.Add(1) + _ = request.Body.Close() + return nil, transportError + })}} + request, err := http.NewRequestWithContext(context.Background(), http.MethodPost, "http://synthetic.invalid/", strings.NewReader("{}")) + if err != nil { + t.Fatal("request construction failed") + } + request.GetBody = nil + _, err = b.sendRemoteRequest(request) + if !errors.Is(err, errBrokerAmbiguous) || errors.Is(err, errBrokerRequestUnsent) || attempts.Load() != 1 { + t.Fatal("transport-returned error allowed unsent classification or replay") + } + }) + } +} diff --git a/broker_zero_invocation_test.go b/broker_zero_invocation_test.go new file mode 100644 index 0000000..bb0e477 --- /dev/null +++ b/broker_zero_invocation_test.go @@ -0,0 +1,127 @@ +package main + +import ( + "bytes" + "context" + "encoding/json" + "io" + "net/http" + "net/http/httptest" + "strings" + "sync/atomic" + "testing" + "time" +) + +func TestBrokerUnsentPromptPreservesRemoteConversation(t *testing.T) { + for _, mode := range []string{"no-invocations", "expiry-before-inference", "reserved-on-restart", "rejected"} { + t.Run(mode, func(t *testing.T) { + f := newBrokerFixture(t, "success") + f.server.Close() + var historyLost, rejectNext atomic.Bool + f.server = httptest.NewServer(http.HandlerFunc(func(w http.ResponseWriter, r *http.Request) { + if r.Method == http.MethodPost && strings.HasSuffix(r.URL.Path, ":stop") { + historyLost.Store(true) + } + if r.Method == http.MethodPost && strings.HasSuffix(r.URL.Path, "/protocols/openai/responses") { + raw, err := io.ReadAll(r.Body) + var request foundryResponseRequest + if err != nil || json.Unmarshal(raw, &request) != nil { + w.WriteHeader(http.StatusBadRequest) + return + } + r.Body = io.NopCloser(bytes.NewReader(raw)) + if rejectNext.Swap(false) { + w.WriteHeader(http.StatusTooManyRequests) + return + } + if request.PreviousResponseID != "" && historyLost.Load() { + w.WriteHeader(http.StatusBadRequest) + return + } + } + f.serve(w, r) + })) + cfg := brokerTestConfig(t, f) + b, server := startBrokerTest(t, cfg) + c := brokerTestContext(cfg) + status, data, err := brokerTestHTTP(context.Background(), server.URL, brokerResponsesPath, c, brokerTestBody("")) + var first foundryResponse + if err != nil || status != http.StatusOK || json.Unmarshal(data, &first) != nil { + t.Fatal("initial fixture response failed") + } + _ = brokerTestControl(t, server.URL, brokerSettlePath, c) + + c.TaskUID, c.PromptID, c.OperationID = "unsent-task", "unsent-prompt", "unsent-invocation" + if mode == "expiry-before-inference" { + c.LeaseExpiresAt = time.Now().Add(500 * time.Millisecond).UTC().Format(time.RFC3339Nano) + } + _ = brokerTestControl(t, server.URL, brokerRenewPath, c) + switch mode { + case "expiry-before-inference": + brokerAwait(t, func() bool { + b.mu.Lock() + defer b.mu.Unlock() + return b.ledger.Sessions[brokerJSONDigest(c.Owner)].Prompts[c.promptKey()].Settled + }) + case "reserved-on-restart": + // A crash after durable reservation leaves no live request and + // never resumes the reserved invocation on broker restart. + c.BodySHA256 = brokerSHA(brokerTestBody(first.ID)) + b.mu.Lock() + err := b.commitLocked(func(next *brokerLedger) error { + session := next.Sessions[brokerJSONDigest(c.Owner)] + if _, err := brokerRecordOperation(session, brokerResponsesPath, c); err != nil { + return err + } + prompt := session.Prompts[c.promptKey()] + prompt.LastSequence = c.InvocationSequence + prompt.Invocations[c.InvocationSequence] = &brokerInvocation{ + Sequence: c.InvocationSequence, OperationID: c.OperationID, BodyDigest: c.BodySHA256, State: "reserved", + } + return nil + }) + valid := brokerLedgerValid(b.ledger, cfg.configDigest) + b.mu.Unlock() + if err != nil || !valid { + t.Fatal("reserved crash fixture did not preserve valid durable ownership") + } + b.close() + server.Close() + b, server = startBrokerTest(t, cfg) + case "rejected": + rejectNext.Store(true) + status, _, err := brokerTestHTTP(context.Background(), server.URL, brokerResponsesPath, c, brokerTestBody(first.ID)) + if err != nil || status == http.StatusOK { + t.Fatal("definite admission rejection was exposed as a response") + } + } + proof := brokerTestControl(t, server.URL, brokerSettlePath, c) + if !proof.SettlementProven || proof.ActiveInvocations != 0 || proof.AmbiguousInvocations != 0 || !brokerDigestValid(proof.ProofDigest) { + t.Fatal("unsent prompt did not obtain durable settlement") + } + _, _, stops, _ := f.counts() + if stops != 0 { + t.Error("settling an unsent prompt stopped the retained remote conversation") + } + c.TaskUID, c.PromptID, c.OperationID = "continued-task", "continued-prompt", "continued-invocation" + c.LeaseExpiresAt = time.Now().Add(10 * time.Second).UTC().Format(time.RFC3339Nano) + status, _, err = brokerTestHTTP(context.Background(), server.URL, brokerResponsesPath, c, brokerTestBody(first.ID)) + if err != nil || status != http.StatusOK { + t.Fatalf("unsent prompt destroyed previous-response continuation: status=%d", status) + } + _ = brokerTestControl(t, server.URL, brokerSettlePath, c) + proof = brokerTestControl(t, server.URL, brokerRetirePath, c) + creates, inferences, stops, deletes := f.counts() + if !proof.RetirementProven || creates != 1 || inferences != 2 || stops != 0 || deletes != 1 { + t.Fatal("unsent prompt cleanup replayed work or lost retirement ownership") + } + b.mu.Lock() + valid := brokerLedgerValid(b.ledger, cfg.configDigest) + b.mu.Unlock() + if !valid { + t.Fatal("unsent prompt settlement invalidated durable ownership") + } + }) + } +} diff --git a/config.go b/config.go index 855f6e5..f0f363f 100644 --- a/config.go +++ b/config.go @@ -50,6 +50,12 @@ const ( envIsolationMode = "ORKA_FOUNDRY_ISOLATION_MODE" envFoundryFeatures = "ORKA_FOUNDRY_FEATURES" envBrokeredToolClasses = "ORKA_FOUNDRY_BROKERED_TOOL_CLASSES" + envToolSchemaMode = "ORKA_FOUNDRY_TOOL_SCHEMA_MODE" +) + +const ( + toolSchemaModeRequest = "request" + toolSchemaModeProviderStatic = "provider-static" ) const foundryEndpointRequirement = "Foundry endpoint must use https " + @@ -82,6 +88,7 @@ type config struct { maxConcurrent int brokeredToolClasses []harness.BrokeredToolClass brokeredToolClassSetting string + toolSchemaMode string } func loadConfig() config { @@ -112,6 +119,7 @@ func loadConfig() config { maxConcurrent: defaultMaxTurns, brokeredToolClasses: brokeredToolClasses, brokeredToolClassSetting: brokeredToolClassSetting, + toolSchemaMode: strings.ToLower(firstNonBlank(os.Getenv(envToolSchemaMode), toolSchemaModeRequest)), } } @@ -191,6 +199,11 @@ func (c config) validate() error { return fmt.Errorf("unsupported Foundry brokered tool class %q", value) } } + switch strings.ToLower(strings.TrimSpace(c.toolSchemaMode)) { + case "", toolSchemaModeRequest, toolSchemaModeProviderStatic: + default: + return fmt.Errorf("foundry tool schema mode must be %s or %s", toolSchemaModeRequest, toolSchemaModeProviderStatic) + } switch strings.ToLower(strings.TrimSpace(c.isolationMode)) { case "entra", "header": default: diff --git a/config_test.go b/config_test.go index a6221c9..c8afd27 100644 --- a/config_test.go +++ b/config_test.go @@ -57,6 +57,20 @@ func TestDefaultFoundryFeaturesHonorsExplicitEmptyValue(t *testing.T) { } } +func TestToolSchemaModeConfiguration(t *testing.T) { + t.Setenv(envToolSchemaMode, toolSchemaModeProviderStatic) + if got := loadConfig().toolSchemaMode; got != toolSchemaModeProviderStatic { + t.Fatalf("tool schema mode = %q, want %q", got, toolSchemaModeProviderStatic) + } + + cfg := testConfig("https://account.services.ai.azure.com") + cfg.projectEndpoint = "https://account.services.ai.azure.com/api/projects/demo" + cfg.toolSchemaMode = "unsupported" + if err := cfg.validate(); err == nil { + t.Fatal("unsupported tool schema mode was accepted") + } +} + func TestConfigRejectsCrossOriginOrWrongAgentResponsesEndpoint(t *testing.T) { base := testConfig("https://account.services.ai.azure.com") base.adapterBearer = "adapter-token" diff --git a/docs/foundry-hosted-v2.md b/docs/foundry-hosted-v2.md new file mode 100644 index 0000000..213c389 --- /dev/null +++ b/docs/foundry-hosted-v2.md @@ -0,0 +1,169 @@ +# Run the v2 supervisor in Foundry + +`Dockerfile.hosted` packages this adapter and Orka's Linux supervisor as a +Foundry Hosted Agent. Orka reaches it through a single-replica gateway in +Kubernetes. The [Foundry lifecycle broker](harness-v2.md) runs beside that +gateway and owns the downstream Responses agent's sessions. + +```mermaid +flowchart LR + Orka -->|harness v2| Gateway + Gateway <-->|two authenticated WebSockets| Hosted[Foundry: launcher + supervisor + ACP child] + Hosted -->|reverse channel| Gateway + Gateway --> Broker + Gateway -->|governed tools and artifacts| Orka + Broker --> Responses[Foundry Responses agent] +``` + +The hosted launcher exposes `GET /readiness` and `/invocations_ws`. Foundry's +endpoint must advertise `invocations_ws` version `2.0.0` and use Entra +authorization. The gateway carries HTTP/2 inside the two WebSockets, preserving +v2 authentication, operation capabilities, fences, streaming responses and +HTTP status codes. It does not retry or reconnect after a channel failure. + +## Build + +First build the configured ACP image and Orka composition as described in +[harness-v2.md](harness-v2.md). The baked `/agent/foundry.json` selects the +downstream Responses agent, model, and tool schema mode. + +Create a public `hosted.json` with these fields: + +| Field | Value | +| --- | --- | +| `protocol` | `orka.foundry.hosted.v1` (the transport bootstrap protocol). | +| `deploymentID` | A new canonical UUID for this deployment. | +| `target` | `projectEndpoint`, `agentName`, and concrete `agentVersion` of the **supervisor's** Hosted Agent. | +| `signingPublicKey` | Ed25519 public key, unpadded URL-safe base64, exactly 32 decoded bytes. | +| `agentConfigurationDigest` | SHA-256 of the exact `/agent/foundry.json` bytes, prefixed with `sha256:`. | + +Keep the signing private key outside every build context. Only the public +configuration belongs in the hosted image. Create a new deployment identity +and key when replacing a hosted lifetime. + +```sh +docker buildx build --builder remote-vm --platform linux/amd64 \ + -f Dockerfile.hosted \ + --build-arg ORKA_RUNTIME_IMAGE=@sha256: \ + --build-arg FOUNDRY_ADAPTER_DIGEST=sha256: \ + --build-arg HOSTED_CONFIG= \ + --provenance=false --push -t /foundry-hosted: . +``` + +Publish the resulting digest as the exact Foundry agent version in `target`. +Use Linux amd64, `PORT=8088`, and root for the launcher and supervisor. The +supervisor assigns a distinct UID/GID to each ACP child. The image uses the +filesystem exercised by the hosted platform probe; a distroless filesystem +failed to start in that environment. + +## Gateway + +Run the configured ACP image in Kubernetes with: + +```text +/agent-runtime-foundry --protocol hosted-gateway --config /etc/orka-foundry/gateway.json +``` + +The public gateway configuration contains: + +| Field | Value | +| --- | --- | +| `protocol` | `orka.foundry.hosted.v1`. | +| `image` | The complete public `hosted.json` object baked into the hosted image. | +| `containerImage` | Exact digest-pinned image in the Foundry version. | +| `sessionID` | A new canonical UUID, chosen before any session creation. | +| `runtimeProfileDigest` | Expected canonical Orka v2 profile digest. | +| `runtimeEnvironment` | The fixed supervisor profile and fence settings below. | +| `orkaBaseURL` | Fixed Orka API origin, for example `http://orka-api.orka-system.svc:8080`. | +| `brokerBaseURL` | Loopback broker origin, for example `http://127.0.0.1:8091`. | + +`runtimeEnvironment` requires `ORKA_ACP_PROVIDER=foundry`, model, adapter digest, +agent configuration digest, tool/approval/MCP policy digests, workspace intent, +`ORKA_ACP_PROXY_CREDENTIAL_ROLE=operator-managed`, +`ORKA_ACP_PROXY_CREDENTIAL_SCOPE=external-runtime`, +`ORKA_ACP_RESOURCE_CLASS=external`, trust namespace, controller epoch, runtime +pool UUID and generation. These use the supervisor variable names documented +by Orka. Optional model context/output limits must be supplied together. +The launcher controls addresses, session directories, boot identity and +credentials; they cannot be overridden through this map. + +Mount these gateway-only files from Kubernetes Secrets: + +| Gateway variable | File contents | +| --- | --- | +| `ORKA_FOUNDRY_GATEWAY_SIGNING_KEY_FILE` | Unpadded URL-safe base64 of the 64-byte Go Ed25519 private key (seed followed by public key), without a newline. | +| `ORKA_FOUNDRY_GATEWAY_CONTROLLER_TOKEN_FILE` | Orka controller bearer, at least 32 bytes. | +| `ORKA_FOUNDRY_GATEWAY_CAPABILITY_SECRET_FILE` | Orka operation capability signing secret, at least 32 bytes. | +| `ORKA_FOUNDRY_GATEWAY_PROVIDER_TOKEN_FILE` | Bearer for the adjacent Foundry broker, at least 32 bytes. | + +Set `ORKA_FOUNDRY_GATEWAY_STATE_DIR` to a private `0700` subdirectory of a +persistent volume. `ORKA_FOUNDRY_GATEWAY_ADDR` defaults to `:8080`. +Give the gateway and broker refreshable Azure Workload Identity. The broker +uses a separate persistent directory and the same baked `foundry.json`. +Use one replica and `Recreate`; never share either ledger between active +writers. `/healthz` becomes ready after bootstrap and profile verification. + +Register the gateway Service as Orka's external `orka.harness.v2` AgentRuntime. +Use the configured ACP image digest as `foundry-serve-acp`'s adapter digest. +Do not set Kubernetes supervisor-recovery metadata: the supervised process +runs in Foundry. Orka's operation credentials still protect every status or +mutation request. Only health and capabilities are safe unauthenticated probes. + +## Ownership and failure behavior + +The gateway verifies the immutable image/version and exact session, and pins +the Azure caller identity. It authenticates both channel roles against one +boot nonce and bootstrap digest. Before sending credentials it fsyncs possible +delivery in its ledger. A gateway restart after that point refuses to seed +another supervisor, even when the earlier acknowledgment was lost. + +Loss of either channel, supervisor exit, Foundry stop/resume, or a container +replacement closes the hosted lifetime. No automatic respawn or session +adoption occurs. In-flight operations may have an unknown outcome. Preserve +both ledgers and use confirmed Orka drain/retirement before replacing a +surviving runtime. After supervisor loss, the current v2 contract cannot +import the old broker's retirement proof; unresolved work remains +`OutcomeUnknown`. A new session ID does not establish that old work stopped. + +The broker settles fully completed foreground Responses from their validated, +durable completion records. It keeps the downstream session running so a +backend can retain conversation history in memory. Interrupted requests still +require remote containment, and session retirement requires confirmed deletion. +Session creation already recorded as an intent gets one bounded attempt that +survives prompt cancellation, so its acknowledgment can be retained without +submitting inference for a cancelled prompt. A lost acknowledgment remains +unresolved ownership. + +Microsoft documents an approximately **30-minute maximum connection duration** +for `invocations_ws`, after which Foundry closes the WebSocket with code `1001`. +This connection limit is separate from the session idle timeout; recent Tasks, +tool calls, or keepalive traffic do not establish that a connection can outlive +it. See [Maximum connection duration](https://learn.microsoft.com/en-us/azure/foundry/agents/how-to/build-voice-agent#maximum-connection-duration). + +Both channel roles belong to one hosted lifetime. A platform close on either +role makes the gateway unavailable and cancels and joins supervisor shutdown. +The gateway does not reconnect, replay an uncertain operation, or reseed the +same ledger after restart. Plan drain and confirmed retirement before the +connection limit, allowing time for cleanup; this package does not support an +indefinitely running hosted supervisor. No application keepalive is sent over +idle channels, so an idle timeout can also end the lifetime earlier. + +Each channel emits one close summary containing only its role, local WebSocket +upgrade time (`opened_at`), elapsed milliseconds (`duration_ms`), and numeric +`close_code`. The code is `0` when no close code was observed before local +closure; `1006` can indicate abnormal EOF without a received close frame. +Peer close text, raw errors, endpoint URLs, and credentials are not logged. +These observations help distinguish a received platform `1001` from an +unclassified local close; elapsed time alone does not prove the cause. + +Bootstrap secrets and supervisor session files stay outside Foundry `HOME`, +which is exposed through its Session Files API. ACP children inherit neither +Azure identity nor gateway/broker/controller credentials. Reverse requests +preserve their original authorization; the gateway never adds a broker token +based on unsigned context. Only the fixed broker lifecycle and Orka MCP/artifact +routes are reachable through the reverse channel. + +The existing Foundry ACP limitations still apply: text/resource-link input, +brokered tools, no per-Task configuration, and no approval-required tools. +Use the Orka harness v2 conformance suite for v2 claims; this repository's +`conformance/` package tests harness v1. diff --git a/docs/harness-v2.md b/docs/harness-v2.md new file mode 100644 index 0000000..45aa83b --- /dev/null +++ b/docs/harness-v2.md @@ -0,0 +1,150 @@ +# Orka harness v2 + +The v2 adapter is an ACP child of Orka's existing supervisor. A separate +Foundry broker owns Azure identity, remote session creation, and durable +cleanup records. The ACP child has no Azure credentials. + +Use one Pod with a supervisor container and a broker sidecar, one replica, +and the Recreate strategy. The broker listens on `127.0.0.1:8091`. Its state +directory must be on a persistent volume that survives either container or +Pod replacement. Mount that volume only in the broker. Never share the +broker's Azure identity or durable state with the ACP child. + +## Build and register + +Create a nonsecret JSON configuration using +[`examples/foundry-acp.json`](../examples/foundry-acp.json). Pin a concrete +Hosted Agent version. Both processes verify the SHA-256 of the exact file +bytes. The target project, agent, version, model, and schema mode therefore +belong to the registered runtime profile. Per-Task configuration overrides +are rejected. + +Build the configured Foundry source image: + +```sh +docker buildx build --platform linux/amd64 -f Dockerfile.acp \ + --build-arg FOUNDRY_CONFIG=examples/foundry-acp.json \ + -t /foundry-configured: --push . +``` + +In the Orka checkout, compose the supervisor with that image: + +```sh +make docker-build-acp-foundry-runtime \ + FOUNDRY_RUNTIME_IMAGE=/foundry-configured@sha256: \ + FOUNDRY_ADAPTER_DIGEST=sha256: \ + ACP_FOUNDRY_RUNTIME_IMG=/foundry-runtime: +``` + +Register the composed image's service as an external `orka.harness.v2` +AgentRuntime. Set `providerKind: foundry`, `adapterName: foundry-serve-acp`, +and `adapterDigest` to the configured Foundry source image digest. The +`agentConfigurationDigest` is `sha256:` plus the SHA-256 of `/agent/foundry.json`. +Set `supportsAgentSessionConfiguration: false`. Follow Orka's external v2 +registration contract for the remaining profile, operation authentication, +controller epoch, workspace, and MCP settings. + +The child accepts text prompts and resource links represented as text. It +uses the supervisor's HTTP MCP server for tools. It does not support image, +audio, embedded-resource, permission, terminal, filesystem, or session-load +ACP requests. Keep approval-required tools empty. The Task must use the +runtime profile's exact brokered tool allowlist. + +The hosted agent's static function names must exactly match those MCP tools. +Use Orka's built-in names or the names of its Kubernetes Tool resources. The +child rejects unknown function names before calling a tool. +Tool results sent to Foundry contain only validated text content, +`structuredContent`, and `isError`. MCP metadata and unrecognized extension +fields are not forwarded. + +## Process configuration + +The supervisor starts the child as: + +```text +/agent-runtime-foundry --protocol acp --config /agent/foundry.json +``` + +Orka supplies these child-only values: + +| Variable | Value | +| --- | --- | +| `ORKA_FOUNDRY_ACP_PROVIDER_BASE_URL` | Per-session supervisor loopback proxy. | +| `ORKA_FOUNDRY_ACP_PROVIDER_TOKEN` | Ephemeral local proxy credential. | +| `ORKA_FOUNDRY_ACP_MODEL` | Model in the baked JSON. | +| `ORKA_FOUNDRY_ACP_AGENT_CONFIGURATION_DIGEST` | Exact baked-file SHA-256. | + +Run the configured source image as the broker sidecar with arguments +`--protocol broker --config /agent/foundry.json`. Give it the same model and +configuration digest, plus: + +| Variable | Value | +| --- | --- | +| `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. | + +Only the broker receives Azure Workload Identity or another refreshable +`DefaultAzureCredential` configuration. The initial implementation requires +Entra isolation. The supervisor's `ORKA_ACP_PROVIDER_PROXY_BASE_URL` points to +`http://127.0.0.1:8091/v1`; its provider token file contains the broker bearer. +Neither the broker bearer nor Azure identity enters the child environment. + +For a Kubernetes exec readiness probe, run: + +```text +/agent-runtime-foundry --protocol broker --health-check +``` + +This checks the configured loopback `/healthz` without loading Azure identity +or opening the ledger. Set the state directory to a private subdirectory of +the persistent volume, such as `/broker-state/ledger`, that the broker user +can create with mode `0700`. + +The broker's persistent directory is private to its OS user. Its ledger +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. + +## Lifecycle guarantees and limits + +The broker persists a random caller-chosen remote session ID before sending +create. Every inference request binds the exact Orka runtime fence, Task, +attempt, prompt digest, lease, invocation sequence, and request-body digest. +It never repeats an inference request after ambiguous acceptance. Remote +response IDs are represented to the child by owner-scoped opaque aliases. + +Successful continuation retains only the last successful response alias. +Failed or cancelled prompts do not become successful conversation history. +`request` mode sends the discovered function schemas to Foundry. The configured +Hosted Agent version's Responses endpoint must accept top-level function tools; +container readiness and SDK support do not prove that Foundry ingress accepts +them. If that endpoint rejects request-level `tools`, use `provider-static` and +preconfigure matching schemas in the Hosted Agent. This mode omits request-level +schemas while retaining the current MCP allowlist. + +Prompt completion and cancellation require remote settlement proof. Session +deletion also requires remote retirement proof. A closed HTTP connection, +local process death, or a single `404` is insufficient when a create or +inference request has an unresolved acceptance outcome. Those cases remain +blocked and Orka reports `OutcomeUnknown` without replaying the prompt. + +Lease expiry and broker startup trigger cleanup of exactly owned sessions. +Kubernetes container-termination recovery does not apply to Foundry because +remote execution can survive the local container. Preserve unresolved +ownership records for investigation; do not fabricate retirement receipts +or remove finalizers to bypass them. + +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 +blocked even if the broker later contains the remote work. + +## Verification + +Run `make verify`. The ACP tests cover real stdio pipes, loopback provider and +MCP servers, continuation, cancellation, malformed Responses streams, +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. diff --git a/examples/foundry-acp.json b/examples/foundry-acp.json new file mode 100644 index 0000000..3ab6845 --- /dev/null +++ b/examples/foundry-acp.json @@ -0,0 +1,9 @@ +{ + "model": "gpt-4.1", + "toolSchemaMode": "provider-static", + "hostedTarget": { + "projectEndpoint": "https://account.services.ai.azure.com/api/projects/project", + "agentName": "incident-scout", + "agentVersion": "2" + } +} diff --git a/go.mod b/go.mod index bafa89f..32543b3 100644 --- a/go.mod +++ b/go.mod @@ -3,9 +3,12 @@ module github.com/orka-agents/agent-runtime-foundry go 1.26.5 require ( + filippo.io/edwards25519 v1.2.0 github.com/Azure/azure-sdk-for-go/sdk/azcore v1.22.0 github.com/Azure/azure-sdk-for-go/sdk/azidentity v1.14.0 github.com/google/uuid v1.6.0 + github.com/gorilla/websocket v1.5.3 + golang.org/x/net v0.55.0 ) require ( @@ -15,7 +18,6 @@ require ( github.com/kylelemons/godebug v1.1.0 // indirect github.com/pkg/browser v0.0.0-20240102092130-5ac0b6a4141c // 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 ) diff --git a/go.sum b/go.sum index c61dd65..3055e1d 100644 --- a/go.sum +++ b/go.sum @@ -1,3 +1,5 @@ +filippo.io/edwards25519 v1.2.0 h1:crnVqOiS4jqYleHd9vaKZ+HKtHfllngJIiOpNpoJsjo= +filippo.io/edwards25519 v1.2.0/go.mod h1:xzAOLCNug/yB62zG1bQ8uziwrIqIuxhctzJT18Q77mc= github.com/Azure/azure-sdk-for-go/sdk/azcore v1.22.0 h1:aokoqcHvaGjiM3VpjKDfMMnF/8epJ+Q1HLJ7CudztqE= github.com/Azure/azure-sdk-for-go/sdk/azcore v1.22.0/go.mod h1:/WYEx9pcM9Y+Dd/APJaNlSvVSvzl54rrMdZT5+Oi2LM= github.com/Azure/azure-sdk-for-go/sdk/azidentity v1.14.0 h1:CU4+EJeJi3TKYWEcYuSdWsjzw0nVsK/H0MSQOiPcymU= @@ -16,6 +18,8 @@ github.com/golang-jwt/jwt/v5 v5.3.1 h1:kYf81DTWFe7t+1VvL7eS+jKFVWaUnK9cB1qbwn63Y github.com/golang-jwt/jwt/v5 v5.3.1/go.mod h1:fxCRLWMO43lRc8nhHWY6LGqRcf+1gQWArsqaEUEa5bE= 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/gorilla/websocket v1.5.3 h1:saDtZ6Pbx/0u+bgYQ3q96pZgCzfhKXGPqt7kZ72aNNg= +github.com/gorilla/websocket v1.5.3/go.mod h1:YR8l580nyteQvAITg2hZ9XVh4b55+EU/adAjf1fMHhE= github.com/keybase/go-keychain v0.0.1 h1:way+bWYa6lDppZoZcgMbYsvC7GxljxrskdNInRtuthU= github.com/keybase/go-keychain v0.0.1/go.mod h1:PdEILRW3i9D8JcdM+FmY6RwkHGnhHxXwkPPMeUgOK1k= github.com/kylelemons/godebug v1.1.0 h1:RPNrshWIDI6G2gRW9EHilWtl7Z6Sb1BR0xunSBf0SNc= diff --git a/hosted_boundary_test.go b/hosted_boundary_test.go new file mode 100644 index 0000000..e12189d --- /dev/null +++ b/hosted_boundary_test.go @@ -0,0 +1,582 @@ +package main + +import ( + "bufio" + "bytes" + "compress/gzip" + "context" + "encoding/base64" + "encoding/json" + "errors" + "io" + "maps" + "net/http" + "net/http/httptest" + "net/url" + "os" + "path/filepath" + "reflect" + "strconv" + "strings" + "sync/atomic" + "testing" + "time" + + "github.com/google/uuid" +) + +func hostedBoundaryLedgerFixture(t *testing.T) (hostedGatewayConfig, hostedGatewayLedger) { + t.Helper() + f := newHostedProtocolFixture(t) + cfg := hostedGatewayConfig{ + Protocol: hostedProtocol, Image: f.config, + ContainerImage: "example.invalid/hosted@" + brokerSHA([]byte("fixture image")), + SessionID: f.hello.Challenge.SessionID, RuntimeProfileDigest: brokerSHA([]byte("fixture profile")), + RuntimeEnvironment: maps.Clone(f.bootstrap.Environment), + OrkaBaseURL: "http://orka.test:8080", BrokerBaseURL: "http://127.0.0.1:8091", + } + ledger := hostedGatewayLedger{ + Version: 1, ConfigDigest: brokerJSONDigest(cfg), SessionID: cfg.SessionID, + PrincipalDigest: brokerSHA([]byte("fixture principal")), CreateAttempted: true, SessionCreated: true, + ExposurePossible: true, Challenge: f.hello.Challenge, PairID: f.hello.PairID, + BootstrapDigest: f.hello.BootstrapDigest, + } + if validateHostedGatewayConfig(cfg) != nil || !hostedGatewayLedgerValid(ledger, cfg) { + t.Fatal("invalid hosted boundary fixture") + } + return cfg, ledger +} + +func hostedBoundaryWriteLedger(t *testing.T, data []byte) string { + t.Helper() + dir := filepath.Join(t.TempDir(), "gateway") + if os.Mkdir(dir, 0o700) != nil || os.WriteFile(filepath.Join(dir, "state.json"), data, 0o600) != nil { + t.Fatal("could not prepare private ledger fixture") + } + return dir +} + +func TestHostedBoundaryLedgerRejectsMixedAndMalformedExposure(t *testing.T) { + cfg, baseline := hostedBoundaryLedgerFixture(t) + for name, mutate := range map[string]func(*hostedGatewayLedger){ + "version": func(v *hostedGatewayLedger) { v.Version++ }, + "config digest": func(v *hostedGatewayLedger) { v.ConfigDigest = brokerSHA([]byte("other config")) }, + "logical session": func(v *hostedGatewayLedger) { v.SessionID = uuid.NewString() }, + "missing principal": func(v *hostedGatewayLedger) { v.PrincipalDigest = "" }, + "malformed principal": func(v *hostedGatewayLedger) { v.PrincipalDigest = "invalid" }, + "unattempted creation": func(v *hostedGatewayLedger) { v.CreateAttempted = false }, + "unexposed creation without principal": func(v *hostedGatewayLedger) { + *v = hostedGatewayLedger{Version: v.Version, ConfigDigest: v.ConfigDigest, SessionID: v.SessionID, CreateAttempted: true} + }, + "uncreated exposure": func(v *hostedGatewayLedger) { v.SessionCreated = false }, + "unrecorded exposure": func(v *hostedGatewayLedger) { v.ExposurePossible = false }, + "missing pair": func(v *hostedGatewayLedger) { v.PairID = "" }, + "nil pair": func(v *hostedGatewayLedger) { v.PairID = uuid.Nil.String() }, + "noncanonical pair": func(v *hostedGatewayLedger) { v.PairID = "AAAAAAAA-AAAA-4AAA-8AAA-AAAAAAAAAAAA" }, + "missing bootstrap": func(v *hostedGatewayLedger) { v.BootstrapDigest = "" }, + "malformed bootstrap": func(v *hostedGatewayLedger) { v.BootstrapDigest = "invalid" }, + "challenge protocol": func(v *hostedGatewayLedger) { v.Challenge.Protocol = "other" }, + "challenge deployment": func(v *hostedGatewayLedger) { v.Challenge.DeploymentID = uuid.NewString() }, + "challenge config": func(v *hostedGatewayLedger) { + v.Challenge.ConfigurationDigest = brokerSHA([]byte("other image config")) + }, + "challenge agent": func(v *hostedGatewayLedger) { v.Challenge.AgentName = "other-agent" }, + "challenge version": func(v *hostedGatewayLedger) { v.Challenge.AgentVersion = "9" }, + "challenge session": func(v *hostedGatewayLedger) { v.Challenge.SessionID = uuid.NewString() }, + "missing boot": func(v *hostedGatewayLedger) { v.Challenge.BootID = "" }, + "nil boot": func(v *hostedGatewayLedger) { v.Challenge.BootID = uuid.Nil.String() }, + "noncanonical boot": func(v *hostedGatewayLedger) { v.Challenge.BootID = "AAAAAAAA-AAAA-4AAA-8AAA-AAAAAAAAAAAA" }, + "missing nonce": func(v *hostedGatewayLedger) { v.Challenge.Nonce = "" }, + "padded nonce": func(v *hostedGatewayLedger) { v.Challenge.Nonce += "=" }, + "zero nonce": func(v *hostedGatewayLedger) { + v.Challenge.Nonce = base64.RawURLEncoding.EncodeToString(make([]byte, 32)) + }, + } { + t.Run(name, func(t *testing.T) { + ledger := baseline + mutate(&ledger) + data, err := json.Marshal(ledger) + if err != nil { + t.Fatal("could not serialize invalid-state fixture") + } + dir := hostedBoundaryWriteLedger(t, data) + store, _, err := openHostedGatewayStore(dir, cfg) + if err == nil { + store.close() + t.Fatal("malformed exposure was accepted for recovery") + } + after, readErr := os.ReadFile(filepath.Join(dir, "state.json")) + if readErr != nil || !bytes.Equal(after, data) { + t.Fatal("rejected ledger was changed or discarded") + } + }) + } +} + +func TestHostedBoundaryLedgerRejectsMalformedEncodingAndConfigChanges(t *testing.T) { + cfg, ledger := hostedBoundaryLedgerFixture(t) + baseline, err := json.Marshal(ledger) + if err != nil { + t.Fatal("could not serialize fixture") + } + malformed := map[string][]byte{ + "empty object": []byte("{}"), + "null": []byte("null"), + "truncated": baseline[:len(baseline)-1], + "duplicate field": append([]byte(`{"version":1,`), baseline[1:]...), + "unknown field": append([]byte(`{"invented":true,`), baseline[1:]...), + "trailing object": append(bytes.Clone(baseline), []byte("{}")...), + "too large": bytes.Repeat([]byte(" "), hostedMaxHandshakeBytes+1), + } + for _, field := range []string{"version", "configDigest", "sessionID"} { + var fields map[string]json.RawMessage + if json.Unmarshal(baseline, &fields) != nil { + t.Fatal("could not decode baseline fixture") + } + delete(fields, field) + malformed["missing "+field], err = json.Marshal(fields) + if err != nil { + t.Fatal("could not encode missing identity fixture") + } + fields[field] = json.RawMessage("null") + malformed["null "+field], err = json.Marshal(fields) + if err != nil { + t.Fatal("could not encode null identity fixture") + } + } + for name, data := range malformed { + t.Run(name, func(t *testing.T) { + store, _, err := openHostedGatewayStore(hostedBoundaryWriteLedger(t, data), cfg) + if err == nil { + store.close() + t.Fatal("malformed ledger encoding was accepted") + } + }) + } + for name, mutate := range map[string]func(*hostedGatewayConfig){ + "session": func(c *hostedGatewayConfig) { c.SessionID = uuid.NewString() }, + "profile": func(c *hostedGatewayConfig) { c.RuntimeProfileDigest = brokerSHA([]byte("other profile")) }, + "image": func(c *hostedGatewayConfig) { + c.ContainerImage = "example.invalid/hosted@" + brokerSHA([]byte("other image")) + }, + "target": func(c *hostedGatewayConfig) { c.Image.Target.AgentVersion = "9" }, + "destination": func(c *hostedGatewayConfig) { c.OrkaBaseURL = "http://other.test:8080" }, + "epoch": func(c *hostedGatewayConfig) { c.RuntimeEnvironment["ORKA_ACP_CONTROLLER_EPOCH"] = "2" }, + } { + t.Run(name, func(t *testing.T) { + changed := cfg + changed.RuntimeEnvironment = maps.Clone(cfg.RuntimeEnvironment) + mutate(&changed) + store, _, err := openHostedGatewayStore(hostedBoundaryWriteLedger(t, baseline), changed) + if err == nil { + store.close() + t.Fatal("existing gateway ownership was rebound to different configuration") + } + }) + } +} + +func TestHostedBoundaryLedgerPersistsExposureWithExclusivePrivateOwnership(t *testing.T) { + cfg, exposed := hostedBoundaryLedgerFixture(t) + dir := filepath.Join(t.TempDir(), "gateway") + store, initial, err := openHostedGatewayStore(dir, cfg) + if err != nil || initial.ExposurePossible || initial.CreateAttempted || initial.SessionCreated || initial.Ready || initial.Closed { + t.Fatal("new ledger was not empty") + } + defer store.close() + other, _, err := openHostedGatewayStore(dir, cfg) + if err == nil { + other.close() + t.Fatal("two gateway writers acquired the same ledger") + } + oldFile, err := os.Open(filepath.Join(dir, "state.json")) + if err != nil { + t.Fatal("could not pin prior state inode") + } + defer oldFile.Close() + initialBytes, err := io.ReadAll(oldFile) + if err != nil { + t.Fatal("could not read prior state") + } + if store.save(exposed) != nil { + t.Fatal("exposure could not be persisted") + } + if _, err := oldFile.Seek(0, io.SeekStart); err != nil { + t.Fatal("could not reread prior inode") + } + oldBytes, err := io.ReadAll(oldFile) + if err != nil || !bytes.Equal(oldBytes, initialBytes) { + t.Fatal("save overwrote the existing ledger inode instead of replacing it") + } + for _, name := range []string{"state.json", "gateway.lock"} { + info, err := os.Lstat(filepath.Join(dir, name)) + if err != nil || !brokerPrivateFile(info) { + t.Fatal("gateway ownership file is not private, singly linked and locally owned") + } + } + entries, err := os.ReadDir(dir) + if err != nil || len(entries) != 2 { + t.Fatal("save left temporary ownership files") + } + store.close() + store, restored, err := openHostedGatewayStore(dir, cfg) + if err != nil || !reflect.DeepEqual(restored, exposed) { + t.Fatal("restart lost or changed possible exposure") + } + defer store.close() + restored.Ready, restored.Closed = true, true + if store.save(restored) != nil { + t.Fatal("closed ownership could not be persisted") + } + store.close() + store, terminal, err := openHostedGatewayStore(dir, cfg) + if err != nil || !reflect.DeepEqual(terminal, restored) { + t.Fatal("closed ledger was reset or rebound on restart") + } + store.close() +} + +func TestHostedBoundaryLedgerRejectsNonprivateAndAliasedFiles(t *testing.T) { + cfg, baseline := hostedBoundaryLedgerFixture(t) + data, _ := json.Marshal(baseline) + for _, name := range []string{"directory permissions", "state permissions", "lock permissions", "state symlink", "lock symlink", "state hardlink", "lock hardlink"} { + t.Run(name, func(t *testing.T) { + dir := hostedBoundaryWriteLedger(t, data) + lock := filepath.Join(dir, "gateway.lock") + if os.WriteFile(lock, nil, 0o600) != nil { + t.Fatal("could not create lock fixture") + } + target := filepath.Join(dir, "state.json") + if strings.HasPrefix(name, "lock") { + target = lock + } + var err error + switch { + case name == "directory permissions": + err = os.Chmod(dir, 0o755) + case strings.HasSuffix(name, "permissions"): + err = os.Chmod(target, 0o644) + case strings.HasSuffix(name, "symlink"): + original := target + "-original" + if err = os.Rename(target, original); err == nil { + err = os.Symlink(original, target) + } + case strings.HasSuffix(name, "hardlink"): + err = os.Link(target, target+"-alias") + } + if err != nil { + t.Fatal("could not prepare unsafe file fixture") + } + store, _, err := openHostedGatewayStore(dir, cfg) + if err == nil { + store.close() + t.Fatal("nonprivate or aliased ownership file was accepted") + } + }) + } + for _, dir := range []string{"", ".", "relative/gateway", "/"} { + store, _, err := openHostedGatewayStore(dir, cfg) + if err == nil { + store.close() + t.Fatal("unsafe state directory was accepted") + } + } +} + +func TestHostedBoundaryLedgerSaveFailurePreservesPriorExposure(t *testing.T) { + cfg, exposed := hostedBoundaryLedgerFixture(t) + data, _ := json.Marshal(exposed) + dir := hostedBoundaryWriteLedger(t, data) + store, _, err := openHostedGatewayStore(dir, cfg) + if err != nil { + t.Fatal("could not open exposure fixture") + } + defer store.close() + moved := dir + "-moved" + if os.Rename(dir, moved) != nil { + t.Fatal("could not inject persistence failure") + } + exposed.Ready = true + if store.save(exposed) == nil { + t.Fatal("save succeeded after state directory disappeared") + } + retained, err := os.ReadFile(filepath.Join(moved, "state.json")) + if err != nil || !bytes.Equal(retained, data) { + t.Fatal("failed save changed the prior exposure record") + } +} + +type hostedBoundaryRoundTripper func(*http.Request) (*http.Response, error) + +func (f hostedBoundaryRoundTripper) RoundTrip(r *http.Request) (*http.Response, error) { return f(r) } + +func TestHostedBoundaryProxyRoutes(t *testing.T) { + digest := strings.Repeat("a", 64) + for _, test := range []struct { + name, method, path string + allow func(*http.Request) bool + accepted bool + }{ + {"v2 health", http.MethodGet, "/v2/health", hostedV2Route, true}, + {"v2 prompt", http.MethodPost, "/v2/sessions/session/prompts", hostedV2Route, true}, + {"v2 lease", http.MethodPatch, "/v2/sessions/session", hostedV2Route, true}, + {"v2 delete", http.MethodDelete, "/v2/sessions/session", hostedV2Route, true}, + {"v2 replay query", http.MethodGet, "/v2/prompts/prompt/events?cursor=7", hostedV2Route, true}, + {"v1 denied", http.MethodPost, "/v1/sessions", hostedV2Route, false}, + {"v2 lookalike", http.MethodPost, "/v20/sessions", hostedV2Route, false}, + {"v2 unsupported method", http.MethodTrace, "/v2/health", hostedV2Route, false}, + {"broker responses", http.MethodPost, brokerResponsesPath, hostedBrokerRoute, true}, + {"broker renew", http.MethodPost, brokerRenewPath, hostedBrokerRoute, true}, + {"broker settle", http.MethodPost, brokerSettlePath, hostedBrokerRoute, true}, + {"broker retire", http.MethodPost, brokerRetirePath, hostedBrokerRoute, true}, + {"broker status GET", http.MethodGet, brokerStatusPath, hostedBrokerRoute, true}, + {"broker status POST denied", http.MethodPost, brokerStatusPath, hostedBrokerRoute, false}, + {"broker response GET denied", http.MethodGet, brokerResponsesPath, hostedBrokerRoute, false}, + {"broker query denied", http.MethodPost, brokerResponsesPath + "?target=other", hostedBrokerRoute, false}, + {"broker arbitrary path denied", http.MethodPost, "/v1/other", hostedBrokerRoute, false}, + {"Orka tool", http.MethodPost, "/internal/v2/acp/mcp/tools/call", hostedOrkaRoute, true}, + {"Orka artifact authorization", http.MethodPost, "/internal/v2/acp/artifact-authorizations", hostedOrkaRoute, true}, + {"Orka artifact read", http.MethodGet, "/internal/v2/acp/artifacts/sha256/" + digest, hostedOrkaRoute, true}, + {"Orka artifact write", http.MethodPut, "/internal/v2/acp/artifacts/sha256/" + digest, hostedOrkaRoute, true}, + {"Orka artifact head", http.MethodHead, "/internal/v2/acp/artifacts/sha256/" + digest, hostedOrkaRoute, true}, + {"Orka artifact delete denied", http.MethodDelete, "/internal/v2/acp/artifacts/sha256/" + digest, hostedOrkaRoute, false}, + {"Orka artifact bad digest", http.MethodGet, "/internal/v2/acp/artifacts/sha256/" + digest[:63], hostedOrkaRoute, false}, + {"Orka arbitrary API denied", http.MethodPost, "/api/tasks", hostedOrkaRoute, false}, + {"Orka query denied", http.MethodPost, "/internal/v2/acp/mcp/tools/call?target=other", hostedOrkaRoute, false}, + } { + t.Run(test.name, func(t *testing.T) { + calls := 0 + proxy, err := newHostedProxy("http://fixed.invalid", hostedBoundaryRoundTripper(func(*http.Request) (*http.Response, error) { + calls++ + return &http.Response{StatusCode: http.StatusNoContent, Header: make(http.Header), Body: http.NoBody}, nil + }), test.allow) + if err != nil { + t.Fatal("could not construct proxy fixture") + } + response := httptest.NewRecorder() + proxy.ServeHTTP(response, httptest.NewRequest(test.method, test.path, nil)) + if test.accepted && (calls != 1 || response.Code != http.StatusNoContent) || + !test.accepted && (calls != 0 || response.Code != http.StatusForbidden) { + t.Fatal("route did not enforce the downstream method/path contract") + } + }) + } +} + +func TestHostedBoundaryProxyRejectsAmbiguousPathsBeforeTransport(t *testing.T) { + for name, mutate := range map[string]func(*http.Request){ + "relative": func(r *http.Request) { r.URL.Path = "v2/health" }, + "dot segment": func(r *http.Request) { r.URL.Path = "/v2/../health" }, + "duplicate slash": func(r *http.Request) { r.URL.Path = "/v2//health" }, + "backslash": func(r *http.Request) { r.URL.Path = "/v2/health\\extra" }, + "NUL": func(r *http.Request) { r.URL.Path = "/v2/health\x00" }, + "escaped path": func(r *http.Request) { r.URL.RawPath = "/v2/%68ealth" }, + "fragment": func(r *http.Request) { r.URL.Fragment = "fragment" }, + "userinfo": func(r *http.Request) { r.URL.User = url.User("fixture") }, + "opaque": func(r *http.Request) { r.URL.Opaque = "//other.invalid/v2/health" }, + "oversized path": func(r *http.Request) { r.URL.Path = "/v2/" + strings.Repeat("x", 4096) }, + "oversized query": func(r *http.Request) { r.URL.RawQuery = strings.Repeat("x", 8193) }, + "upgrade": func(r *http.Request) { r.Header.Set("Upgrade", "websocket") }, + "CONNECT": func(r *http.Request) { r.Method = http.MethodConnect }, + "nil URL": func(r *http.Request) { r.URL = nil }, + } { + t.Run(name, func(t *testing.T) { + calls := 0 + proxy, err := newHostedProxy("http://fixed.invalid", hostedBoundaryRoundTripper(func(*http.Request) (*http.Response, error) { + calls++ + return nil, errors.New("disallowed request reached transport") + }), hostedV2Route) + if err != nil { + t.Fatal("could not construct proxy fixture") + } + request := httptest.NewRequest(http.MethodGet, "/v2/health", nil) + mutate(request) + response := httptest.NewRecorder() + proxy.ServeHTTP(response, request) + if calls != 0 || response.Code != http.StatusForbidden { + t.Fatal("ambiguous request reached the fixed destination") + } + }) + } +} + +func TestHostedBoundaryProxyPinsDestinationAndPreservesOnlySuppliedAuthorization(t *testing.T) { + for _, test := range []struct { + name, path string + allow func(*http.Request) bool + }{ + {"broker", brokerRenewPath, hostedBrokerRoute}, + {"Orka", "/internal/v2/acp/mcp/tools/call", hostedOrkaRoute}, + {"supervisor", "/v2/sessions", hostedV2Route}, + } { + for _, authorization := range [][]string{nil, {"Bearer fixture-incoming"}, {"Bearer fixture-first", "Bearer fixture-second"}} { + t.Run(test.name+"-"+strconv.Itoa(len(authorization)), func(t *testing.T) { + calls := 0 + payload := []byte("fixture request") + proxy, err := newHostedProxy("http://fixed.invalid:8091", hostedBoundaryRoundTripper(func(r *http.Request) (*http.Response, error) { + calls++ + body, readErr := io.ReadAll(r.Body) + if readErr != nil || !bytes.Equal(body, payload) || r.URL.Scheme != "http" || + r.URL.Host != "fixed.invalid:8091" || r.Host != "fixed.invalid:8091" || r.URL.Path != test.path || + !reflect.DeepEqual(r.Header.Values("Authorization"), authorization) || r.GetBody != nil { + t.Error("proxy changed authority/body or allowed caller destination selection") + } + if r.Header.Get("Forwarded") != "" || r.Header.Get("X-Forwarded-Host") != "" { + t.Error("proxy retained untrusted forwarded routing headers") + } + return &http.Response{StatusCode: http.StatusAccepted, Header: make(http.Header), Body: http.NoBody}, nil + }), test.allow) + if err != nil { + t.Fatal("could not construct proxy fixture") + } + request := httptest.NewRequest(http.MethodPost, "https://untrusted.invalid"+test.path, bytes.NewReader(payload)) + request.Host = "also-untrusted.invalid" + request.Header["Authorization"] = authorization + request.Header.Set("X-Foundry-Context", "unsigned-fixture-context") + request.Header.Set("Forwarded", "host=untrusted.invalid;proto=https") + request.Header.Set("X-Forwarded-Host", "untrusted.invalid") + request.GetBody = func() (io.ReadCloser, error) { return io.NopCloser(bytes.NewReader(payload)), nil } + response := httptest.NewRecorder() + proxy.ServeHTTP(response, request) + if response.Code != http.StatusAccepted || calls != 1 { + t.Fatal("authorized fixed-route fixture failed or retried") + } + }) + } + } + for _, target := range []string{"", "file:///tmp/socket", "unix:///tmp/socket", "http:///missing-host", "http://fixed.invalid/base"} { + if _, err := newHostedProxy(target, nil, hostedV2Route); err == nil { + t.Fatal("invalid proxy destination was accepted") + } + } + for _, target := range []string{"http://remote.invalid", "http://localhost:8091", "http://10.0.0.1:8091", "http://127.0.0.1:8091/base", "http://127.0.0.1:8091?target=other", "http://user@127.0.0.1:8091"} { + if hostedRelayTargetValid(target, true) { + t.Fatal("broker destination escaped its configured loopback boundary") + } + } +} + +func TestHostedBoundaryProxyPreservesStreamingStatusHeadersAndTrailers(t *testing.T) { + release := make(chan struct{}) + captured := make(chan bool, 1) + backend := httptest.NewServer(http.HandlerFunc(func(w http.ResponseWriter, r *http.Request) { + body, err := io.ReadAll(r.Body) + captured <- err == nil && string(body) == "fixture request" && r.Header.Get("Authorization") == "Bearer fixture-incoming" + w.Header().Set("Content-Type", "application/x-ndjson") + w.Header().Add("X-Orka-Fence", "first") + w.Header().Add("X-Orka-Fence", "second") + w.Header().Set("Trailer", "X-Orka-Complete") + w.WriteHeader(http.StatusAccepted) + _, _ = io.WriteString(w, "{\"sequence\":1}\n") + _ = http.NewResponseController(w).Flush() + select { + case <-release: + _, _ = io.WriteString(w, "{\"sequence\":2}\n") + w.Header().Set("X-Orka-Complete", "yes") + case <-r.Context().Done(): + } + })) + defer backend.Close() + transport := newHostedLocalTransport() + defer transport.CloseIdleConnections() + proxy, err := newHostedProxy(backend.URL, transport, hostedV2Route) + if err != nil { + t.Fatal("could not construct streaming proxy") + } + frontend := httptest.NewServer(proxy) + defer frontend.Close() + ctx, cancel := context.WithTimeout(context.Background(), 5*time.Second) + defer cancel() + request, _ := http.NewRequestWithContext(ctx, http.MethodPost, frontend.URL+"/v2/stream", strings.NewReader("fixture request")) + request.Header.Set("Authorization", "Bearer fixture-incoming") + response, err := frontend.Client().Do(request) + if err != nil { + t.Fatal("streaming request failed") + } + defer response.Body.Close() + reader := bufio.NewReader(response.Body) + line, err := reader.ReadString('\n') + if err != nil || line != "{\"sequence\":1}\n" || response.StatusCode != http.StatusAccepted || + response.Header.Get("Content-Type") != "application/x-ndjson" || + !reflect.DeepEqual(response.Header.Values("X-Orka-Fence"), []string{"first", "second"}) || !<-captured { + t.Fatal("streaming prefix, status, headers or request authority changed") + } + close(release) + rest, err := io.ReadAll(reader) + if err != nil || string(rest) != "{\"sequence\":2}\n" || response.Trailer.Get("X-Orka-Complete") != "yes" { + t.Fatal("streaming suffix or terminal trailer changed") + } +} + +func TestHostedBoundaryProxyPreservesEncodedResponseBytes(t *testing.T) { + var compressed bytes.Buffer + writer := gzip.NewWriter(&compressed) + _, _ = writer.Write([]byte("fixture encoded body")) + if writer.Close() != nil { + t.Fatal("could not prepare encoded response") + } + backend := httptest.NewServer(http.HandlerFunc(func(w http.ResponseWriter, r *http.Request) { + w.Header().Set("Content-Encoding", "gzip") + w.Header().Set("Content-Length", strconv.Itoa(compressed.Len())) + w.Header().Set("Content-Type", "application/octet-stream") + w.WriteHeader(http.StatusUnprocessableEntity) + _, _ = w.Write(compressed.Bytes()) + })) + defer backend.Close() + transport := newHostedLocalTransport() + defer transport.CloseIdleConnections() + proxy, err := newHostedProxy(backend.URL, transport, hostedV2Route) + if err != nil { + t.Fatal("could not construct encoded-response proxy") + } + response := httptest.NewRecorder() + proxy.ServeHTTP(response, httptest.NewRequest(http.MethodGet, "/v2/status", nil)) + if response.Code != http.StatusUnprocessableEntity || response.Header().Get("Content-Encoding") != "gzip" || + response.Header().Get("Content-Length") != strconv.Itoa(compressed.Len()) || !bytes.Equal(response.Body.Bytes(), compressed.Bytes()) { + t.Fatal("local hop transparently altered response encoding or bytes") + } +} + +func TestHostedBoundaryProxyDoesNotFollowRedirectsOrRetryAcceptedLocalRequests(t *testing.T) { + var calls atomic.Int32 + backend := httptest.NewServer(http.HandlerFunc(func(w http.ResponseWriter, r *http.Request) { + calls.Add(1) + if r.URL.Path == "/v2/redirect" { + w.Header().Set("Location", "http://must-not-be-contacted.invalid/v2/other") + w.WriteHeader(http.StatusTemporaryRedirect) + _, _ = io.WriteString(w, "fixture redirect") + return + } + _, _ = io.Copy(io.Discard, r.Body) + conn, _, err := http.NewResponseController(w).Hijack() + if err == nil { + _ = conn.Close() + } + })) + defer backend.Close() + transport := newHostedLocalTransport() + defer transport.CloseIdleConnections() + proxy, err := newHostedProxy(backend.URL, transport, hostedV2Route) + if err != nil { + t.Fatal("could not construct local-failure proxy") + } + redirect := httptest.NewRecorder() + proxy.ServeHTTP(redirect, httptest.NewRequest(http.MethodGet, "/v2/redirect", nil)) + if redirect.Code != http.StatusTemporaryRedirect || redirect.Header().Get("Location") != "http://must-not-be-contacted.invalid/v2/other" || + redirect.Body.String() != "fixture redirect" || calls.Load() != 1 { + t.Fatal("proxy changed or followed a downstream redirect") + } + request := httptest.NewRequest(http.MethodPost, "/v2/mutation", strings.NewReader("fixture mutation")) + request.Header.Set("Idempotency-Key", "fixture") + var replays atomic.Int32 + request.GetBody = func() (io.ReadCloser, error) { + replays.Add(1) + return io.NopCloser(strings.NewReader("fixture mutation")), nil + } + failed := httptest.NewRecorder() + proxy.ServeHTTP(failed, request) + if failed.Code != http.StatusBadGateway || calls.Load() != 2 || replays.Load() != 0 || + !strings.Contains(failed.Body.String(), "acceptance may be unknown") { + t.Fatal("accepted local request was replayed or misclassified after disconnection") + } +} diff --git a/hosted_config.go b/hosted_config.go new file mode 100644 index 0000000..f953ade --- /dev/null +++ b/hosted_config.go @@ -0,0 +1,180 @@ +package main + +import ( + "crypto/ed25519" + "encoding/base64" + "io" + "net" + "net/url" + "os" + "path/filepath" + "strconv" + "strings" +) + +const ( + hostedImageConfigPath = "/agent/hosted.json" + hostedGatewayConfigPath = "/etc/orka-foundry/gateway.json" + hostedMaxHandshakeBytes = 64 << 10 + hostedSupervisorAddr = "127.0.0.1:8080" + hostedBrokerRelayAddr = "127.0.0.1:8091" + hostedOrkaRelayAddr = "127.0.0.1:8092" +) + +// Both configuration files are public. Credentials and the signing key are +// read only from gateway-only Secret mounts, never the hosted image or HOME. +type hostedGatewayConfig struct { + Protocol string `json:"protocol"` + Image hostedImageConfig `json:"image"` + ContainerImage string `json:"containerImage"` + SessionID string `json:"sessionID"` + RuntimeProfileDigest string `json:"runtimeProfileDigest"` + RuntimeEnvironment map[string]string `json:"runtimeEnvironment"` + OrkaBaseURL string `json:"orkaBaseURL"` + BrokerBaseURL string `json:"brokerBaseURL"` +} + +type hostedGatewaySettings struct { + config hostedGatewayConfig + bootstrap hostedBootstrap + signingKey ed25519.PrivateKey + stateDir string + address string +} + +func readHostedConfig(path string, target any) error { + data, err := readHostedFile(path, hostedMaxHandshakeBytes) + if err != nil || acpDecode(data, target, true) != nil { + return errHostedInvalid + } + return nil +} + +func readHostedFile(path string, maximum int64) ([]byte, error) { + if !filepath.IsAbs(path) { + return nil, errHostedInvalid + } + file, err := os.Open(path) + if err != nil { + return nil, errHostedInvalid + } + defer file.Close() //nolint:errcheck + info, err := file.Stat() + if err != nil || !info.Mode().IsRegular() || info.Size() > maximum { + return nil, errHostedInvalid + } + data, err := io.ReadAll(io.LimitReader(file, maximum+1)) + if err != nil || int64(len(data)) > maximum { + return nil, errHostedInvalid + } + return data, nil +} + +func loadHostedImageConfig(path string) (hostedImageConfig, error) { + var cfg hostedImageConfig + if readHostedConfig(path, &cfg) != nil || validateHostedImageConfig(cfg) != nil { + return hostedImageConfig{}, errHostedInvalid + } + return cfg, nil +} + +func loadHostedGatewaySettings(path string, getenv func(string) string) (hostedGatewaySettings, error) { + var settings hostedGatewaySettings + if readHostedConfig(path, &settings.config) != nil || validateHostedGatewayConfig(settings.config) != nil { + return settings, errHostedInvalid + } + settings.stateDir = getenv("ORKA_FOUNDRY_GATEWAY_STATE_DIR") + settings.address = firstNonBlank(getenv("ORKA_FOUNDRY_GATEWAY_ADDR"), ":8080") + if !filepath.IsAbs(settings.stateDir) || filepath.Clean(settings.stateDir) == "/" || !hostedListenAddressValid(settings.address) { + return settings, errHostedInvalid + } + key, err := readHostedFile(getenv("ORKA_FOUNDRY_GATEWAY_SIGNING_KEY_FILE"), 4096) + if err != nil { + return settings, errHostedInvalid + } + defer clear(key) + decoded, err := base64.RawURLEncoding.Strict().DecodeString(string(key)) + if err != nil || len(decoded) != ed25519.PrivateKeySize || base64.RawURLEncoding.EncodeToString(decoded) != string(key) { + clear(decoded) + return settings, errHostedInvalid + } + settings.signingKey = ed25519.PrivateKey(decoded) + // Reject unusable local keys before Azure authentication or session creation. + if !hostedSigningKeyValid(settings.signingKey) || + base64.RawURLEncoding.EncodeToString(settings.signingKey.Public().(ed25519.PublicKey)) != settings.config.Image.SigningPublicKey { + clear(decoded) + return settings, errHostedInvalid + } + settings.bootstrap.Environment = settings.config.RuntimeEnvironment + for _, field := range []struct { + env string + value *string + }{ + {"ORKA_FOUNDRY_GATEWAY_CONTROLLER_TOKEN_FILE", &settings.bootstrap.ControllerToken}, + {"ORKA_FOUNDRY_GATEWAY_CAPABILITY_SECRET_FILE", &settings.bootstrap.CapabilitySecret}, + {"ORKA_FOUNDRY_GATEWAY_PROVIDER_TOKEN_FILE", &settings.bootstrap.ProviderToken}, + } { + data, err := readHostedFile(getenv(field.env), 16<<10) + if err != nil { + clear(settings.signingKey) + return settings, errHostedInvalid + } + *field.value = string(data) + clear(data) + } + if validateHostedBootstrap(settings.config.Image, settings.bootstrap) != nil { + clear(settings.signingKey) + return settings, errHostedInvalid + } + return settings, nil +} + +func validateHostedGatewayConfig(cfg hostedGatewayConfig) error { + imageName, digest, pinned := strings.Cut(cfg.ContainerImage, "@") + if cfg.Protocol != hostedProtocol || validateHostedImageConfig(cfg.Image) != nil || + !hostedUUIDValid(cfg.SessionID) || !brokerDigestValid(cfg.RuntimeProfileDigest) || + !pinned || !brokerDigestValid(digest) || !acpSafeString(imageName, 512) || + strings.ContainsAny(imageName, " @\\?#") { + return errHostedInvalid + } + // Validate the environment without reading any credential. These synthetic + // placeholders never leave this validation function. + placeholder := strings.Repeat("x", 32) + if validateHostedBootstrap(cfg.Image, hostedBootstrap{Environment: cfg.RuntimeEnvironment, + ControllerToken: placeholder, CapabilitySecret: placeholder, ProviderToken: placeholder}) != nil { + return errHostedInvalid + } + if !hostedRelayTargetValid(cfg.OrkaBaseURL, false) || !hostedRelayTargetValid(cfg.BrokerBaseURL, true) { + return errHostedInvalid + } + return nil +} + +func hostedRelayTargetValid(raw string, loopback bool) bool { + u, err := url.Parse(raw) + if err != nil || !acpSafeString(raw, 2048) || (u.Scheme != "http" && u.Scheme != "https") || + u.Host == "" || u.User != nil || u.Path != "" || u.RawPath != "" || + u.RawQuery != "" || u.ForceQuery || u.Fragment != "" || strings.Contains(raw, "#") { + return false + } + if port := u.Port(); port != "" { + n, err := strconv.Atoi(port) + if err != nil || n < 1 || n > 65535 { + return false + } + } + if loopback { + ip := net.ParseIP(u.Hostname()) + return ip != nil && ip.IsLoopback() + } + // This is a fixed operator-configured destination, never selected by an + // incoming HTTP request. In-cluster DNS and HTTPS endpoints are supported. + return true +} + +func hostedListenAddressValid(address string) bool { + host, port, err := net.SplitHostPort(address) + n, portErr := strconv.Atoi(port) + return err == nil && portErr == nil && n >= 1 && n <= 65535 && + (host == "" || net.ParseIP(host) != nil) +} diff --git a/hosted_create_rejection_test.go b/hosted_create_rejection_test.go new file mode 100644 index 0000000..d8de3a5 --- /dev/null +++ b/hosted_create_rejection_test.go @@ -0,0 +1,102 @@ +package main + +import ( + "context" + "io" + "net/http" + "path/filepath" + "strconv" + "strings" + "sync/atomic" + "testing" +) + +func TestHostedRejectedCreateAllowsOneLaterStartup(t *testing.T) { + for _, status := range []int{400, 401, 403, 404, 405, 413, 415, 422, 429} { + t.Run(strconv.Itoa(status), func(t *testing.T) { + f := newHostedGatewayTestFixture(t, hostedGatewayTestOptions{}) + transport := f.gateway.httpClient.Transport + var submissions atomic.Int32 + f.gateway.httpClient.Transport = hostedGatewayTestRoundTripper(func(request *http.Request) (*http.Response, error) { + if request.Method == http.MethodPost && submissions.Add(1) == 1 { + return &http.Response{StatusCode: status, Header: make(http.Header), + Body: io.NopCloser(strings.NewReader("{}")), Request: request}, nil + } + return transport.RoundTrip(request) + }) + if f.initialize() == nil || submissions.Load() != 1 || f.creates.Load() != 0 { + t.Fatal("fixture did not reach one complete creation rejection") + } + f.assertNoBootstrap(t) + store, ledger, err := openHostedGatewayStore(f.settings.stateDir, f.settings.config) + if err != nil { + t.Fatal("could not reopen rejected-create ledger") + } + defer store.close() + if ledger.CreateAttempted || ledger.SessionCreated || ledger.ExposurePossible || ledger.PrincipalDigest == "" { + t.Fatal("complete admission rejection retained possibly-sent creation") + } + restarted := &hostedGateway{cfg: f.settings.config, provider: f.gateway.provider, + httpClient: f.gateway.httpClient, store: store, ledger: ledger} + if restarted.ensureSession(t.Context()) != nil || submissions.Load() != 2 || f.creates.Load() != 1 || + !restarted.ledger.SessionCreated || restarted.ledger.PrincipalDigest != ledger.PrincipalDigest { + t.Fatal("later explicit startup did not retain identity and make exactly one creation") + } + }) + } +} + +type hostedIncompleteRejectionBody struct{} + +func (hostedIncompleteRejectionBody) Read([]byte) (int, error) { return 0, io.ErrUnexpectedEOF } +func (hostedIncompleteRejectionBody) Close() error { return nil } + +func TestHostedAmbiguousCreateCannotClearIntent(t *testing.T) { + for _, failure := range []string{"conflict", "server-error", "incomplete-rejection", "rollback-storage"} { + t.Run(failure, func(t *testing.T) { + f := newHostedGatewayTestFixture(t, hostedGatewayTestOptions{}) + transport := f.gateway.httpClient.Transport + var submissions atomic.Int32 + f.gateway.httpClient.Transport = hostedGatewayTestRoundTripper(func(request *http.Request) (*http.Response, error) { + if request.Method != http.MethodPost { + return transport.RoundTrip(request) + } + submissions.Add(1) + status := http.StatusForbidden + var body io.ReadCloser = io.NopCloser(strings.NewReader("{}")) + switch failure { + case "conflict": + status = http.StatusConflict + case "server-error": + status = http.StatusInternalServerError + case "incomplete-rejection": + body = hostedIncompleteRejectionBody{} + case "rollback-storage": + f.gateway.store.dir = filepath.Join(f.settings.stateDir, "missing") + } + return &http.Response{StatusCode: status, Header: make(http.Header), Body: body, Request: request}, nil + }) + if f.initialize() == nil || submissions.Load() != 1 || f.creates.Load() != 0 { + t.Fatal("fixture did not reach the selected submitted-create failure") + } + if !f.gateway.ledger.CreateAttempted || f.gateway.ledger.SessionCreated || f.gateway.ledger.ExposurePossible { + t.Fatal("failed creation cleared in-memory ownership without durable proof") + } + store, ledger, err := openHostedGatewayStore(f.settings.stateDir, f.settings.config) + if err != nil { + t.Fatal("could not reopen submitted-create ledger") + } + defer store.close() + if !ledger.CreateAttempted || ledger.SessionCreated || ledger.ExposurePossible { + t.Fatal("failed creation lost its original durable intent") + } + restarted := &hostedGateway{cfg: f.settings.config, provider: f.gateway.provider, + httpClient: f.gateway.httpClient, store: store, ledger: ledger} + requests, tokens := f.httpCalls.Load(), f.tokenCalls.Load() + if restarted.ensureSession(context.Background()) == nil || submissions.Load() != 1 || + f.httpCalls.Load() != requests || f.tokenCalls.Load() != tokens { + t.Fatal("uncertain creation was retried or adopted") + } + }) + } +} diff --git a/hosted_gateway.go b/hosted_gateway.go new file mode 100644 index 0000000..e47e812 --- /dev/null +++ b/hosted_gateway.go @@ -0,0 +1,245 @@ +package main + +import ( + "context" + "encoding/json" + "io" + "log" + "net" + "net/http" + "sync" + "sync/atomic" + "time" + + "github.com/google/uuid" + "github.com/gorilla/websocket" +) + +type hostedWebSocketDial func(context.Context, string, http.Header) (*websocket.Conn, *http.Response, error) + +type hostedGateway struct { + cfg hostedGatewayConfig + provider foundryTokenProvider + httpClient *http.Client + dial hostedWebSocketDial + store *hostedGatewayStore + ledger hostedGatewayLedger + ctx context.Context + cancel context.CancelFunc + mu sync.Mutex + closeOnce sync.Once + ready atomic.Bool + handler http.Handler + forward, reverse *websocket.Conn + forwardObservation, reverseObservation *hostedChannelObservation + channelLog *log.Logger +} + +func newHostedGateway(ctx context.Context, settings hostedGatewaySettings, provider foundryTokenProvider) (*hostedGateway, error) { + defer clear(settings.signingKey) + if validateHostedGatewayConfig(settings.config) != nil || validateHostedBootstrap(settings.config.Image, settings.bootstrap) != nil || provider == nil { + return nil, errHostedInvalid + } + store, ledger, err := openHostedGatewayStore(settings.stateDir, settings.config) + if err != nil { + return nil, err + } + // This check precedes Azure calls and channel setup, including after a + // process/pod restart. No automatic replacement can inherit unknown work. + if ledger.ExposurePossible { + store.close() + return nil, errHostedInvalid + } + lifetime, cancel := context.WithCancel(ctx) + transport := newHostedLocalTransport() + dialer := &websocket.Dialer{HandshakeTimeout: 120 * time.Second, + NetDialContext: (&net.Dialer{Timeout: 10 * time.Second, KeepAlive: -1}).DialContext, + ReadBufferSize: 4096, WriteBufferSize: 4096} + g := &hostedGateway{cfg: settings.config, provider: provider, store: store, ledger: ledger, + ctx: lifetime, cancel: cancel, dial: dialer.DialContext, + httpClient: &http.Client{Transport: transport, Timeout: 120 * time.Second, + CheckRedirect: func(*http.Request, []*http.Request) error { return errHostedInvalid }}} + if err := g.initialize(settings); err != nil { + g.close() + return nil, err + } + return g, nil +} + +func (g *hostedGateway) initialize(settings hostedGatewaySettings) error { + // JSON escaping can exceed the wire limit even when each raw field is valid. + // Reject that configuration before creating or consuming a hosted lifetime. + body, err := json.Marshal(settings.bootstrap) + defer clear(body) + if err != nil || len(body) > hostedMaxHandshakeBytes { + return errHostedInvalid + } + ctx, cancel := context.WithTimeout(g.ctx, hostedInitializationWait) + defer cancel() + if g.validateRemote(ctx) != nil || g.ensureSession(ctx) != nil { + return errHostedInvalid + } + forward, challenge, err := g.openChannel(ctx, "forward") + if err != nil { + return err + } + g.forward = forward + reverse, otherChallenge, err := g.openChannel(ctx, "reverse") + if err != nil { + return err + } + g.reverse = reverse + if otherChallenge != challenge { + return errHostedInvalid + } + bootstrapDigest, pairID := brokerSHA(body), uuid.NewString() + for _, channel := range []struct { + ws *websocket.Conn + role string + }{{forward, "forward"}, {reverse, "reverse"}} { + hello, err := signHostedHello(hostedHello{Challenge: challenge, PairID: pairID, Role: channel.role, + BootstrapDigest: bootstrapDigest, ExpiresAt: time.Now().Add(60 * time.Second).Unix()}, settings.signingKey) + _ = channel.ws.SetWriteDeadline(time.Now().Add(10 * time.Second)) + if err != nil || channel.ws.WriteJSON(hello) != nil { + return errHostedInvalid + } + var ack hostedAccepted + _ = channel.ws.SetReadDeadline(time.Now().Add(30 * time.Second)) + if readHostedWSJSON(channel.ws, &ack) != nil || ack != (hostedAccepted{Protocol: hostedProtocol, + PairID: pairID, Role: channel.role, BootID: challenge.BootID, BootstrapDigest: bootstrapDigest}) { + return errHostedInvalid + } + } + reverseHandler, err := g.reverseHandler() + if err != nil { + return err + } + reverseConn := newHostedObservedWSConn(reverse, g.reverseObservation) + go serveHostedHTTP2(g.ctx, reverseConn, reverseHandler) + // Fsync this record BEFORE sending even the first bootstrap byte. Failure + // after this point requires explicit retirement and a fresh deployment. + next := g.ledger + next.ExposurePossible, next.Challenge, next.PairID, next.BootstrapDigest = true, challenge, pairID, bootstrapDigest + if g.store.save(next) != nil { + return errHostedInvalid + } + g.ledger = next + _ = forward.SetWriteDeadline(time.Now().Add(10 * time.Second)) + if forward.WriteMessage(websocket.TextMessage, body) != nil { + return errHostedInvalid + } + _ = forward.SetReadDeadline(time.Now().Add(45 * time.Second)) + var ready hostedReady + if readHostedWSJSON(forward, &ready) != nil || ready != (hostedReady{Protocol: hostedProtocol, + PairID: pairID, BootID: challenge.BootID, Ready: true}) { + return errHostedInvalid + } + forwardConn := newHostedObservedWSConn(forward, g.forwardObservation) + client, err := newHostedHTTP2ClientConn(forwardConn) + if err != nil { + return err + } + if g.validateCapabilities(ctx, client) != nil { + return errHostedInvalid + } + handler, err := newHostedProxy("http://supervisor", client, hostedV2Route) + if err != nil { + return err + } + g.handler = handler + next = g.ledger + next.Ready = true + if g.store.save(next) != nil { + return errHostedInvalid + } + g.ledger = next + g.ready.Store(true) + go func() { + select { + case <-g.ctx.Done(): + case <-forwardConn.Done(): + case <-reverseConn.Done(): + } + g.close() + }() + return nil +} + +func (g *hostedGateway) reverseHandler() (http.Handler, error) { + broker, err := newHostedProxy(g.cfg.BrokerBaseURL, newHostedLocalTransport(), hostedBrokerRoute) + if err != nil { + return nil, err + } + orka, err := newHostedProxy(g.cfg.OrkaBaseURL, newHostedLocalTransport(), hostedOrkaRoute) + if err != nil { + return nil, err + } + return http.HandlerFunc(func(w http.ResponseWriter, r *http.Request) { + switch r.Host { + case "broker": + broker.ServeHTTP(w, r) + case "orka": + orka.ServeHTTP(w, r) + default: + http.Error(w, "hosted authority denied", http.StatusForbidden) + } + }), nil +} + +func (g *hostedGateway) validateCapabilities(ctx context.Context, transport http.RoundTripper) error { + request, _ := http.NewRequestWithContext(ctx, http.MethodGet, "http://supervisor/v2/capabilities", nil) + response, err := transport.RoundTrip(request) + if err != nil { + return errHostedInvalid + } + defer response.Body.Close() //nolint:errcheck + data, err := io.ReadAll(io.LimitReader(response.Body, hostedMaxHandshakeBytes+1)) + var value struct { + Protocol string `json:"protocol"` + Transport string `json:"transport"` + RuntimeProfileDigest string `json:"runtimeProfileDigest"` + AdapterDigests map[string]string `json:"adapterDigests"` + } + if err != nil || len(data) > hostedMaxHandshakeBytes || response.StatusCode != http.StatusOK || + acpDecode(data, &value, false) != nil || value.Protocol != "orka.harness.v2" || value.Transport != "http+ndjson" || + value.RuntimeProfileDigest != g.cfg.RuntimeProfileDigest || + value.AdapterDigests["foundry-serve-acp"] != g.cfg.RuntimeEnvironment["ORKA_ACP_FOUNDRY_ADAPTER_DIGEST"] { + return errHostedInvalid + } + return nil +} + +func (g *hostedGateway) ServeHTTP(w http.ResponseWriter, r *http.Request) { + if !g.ready.Load() || g.ctx.Err() != nil { + http.Error(w, "hosted lifetime unavailable; acceptance may be unknown", http.StatusServiceUnavailable) + return + } + if r.Method == http.MethodGet && r.URL.Path == "/healthz" && r.URL.RawQuery == "" { + w.WriteHeader(http.StatusOK) + return + } + g.handler.ServeHTTP(w, r) +} + +func (g *hostedGateway) close() { + g.closeOnce.Do(func() { + g.ready.Store(false) + g.cancel() + if g.forward != nil { + _ = g.forward.Close() + } + if g.reverse != nil { + _ = g.reverse.Close() + } + g.forwardObservation.finish() + g.reverseObservation.finish() + g.mu.Lock() + if g.ledger.ExposurePossible { + g.ledger.Closed = true + _ = g.store.save(g.ledger) // ExposurePossible remains durable even if this write fails. + } + g.store.close() + g.mu.Unlock() + g.httpClient.CloseIdleConnections() + }) +} diff --git a/hosted_gateway_test.go b/hosted_gateway_test.go new file mode 100644 index 0000000..e5eee3d --- /dev/null +++ b/hosted_gateway_test.go @@ -0,0 +1,1024 @@ +package main + +import ( + "bufio" + "bytes" + "context" + "encoding/base64" + "encoding/json" + "errors" + "io" + "maps" + "net" + "net/http" + "net/http/httptest" + "net/url" + "os" + "path/filepath" + "reflect" + "strings" + "sync" + "sync/atomic" + "testing" + "time" + + "github.com/google/uuid" + "github.com/gorilla/websocket" + "golang.org/x/net/http2" +) + +type hostedGatewayTestOptions struct { + remote func(string, map[string]any) + challenge func(int, *hostedChallenge) + accepted func(*hostedAccepted) + ready func(*hostedReady) + capabilities func(map[string]any) + beforeReverseAck func() + supervisor http.Handler + callback http.Handler + preexistingSession bool + ambiguousCreate bool + dropAfterBootstrap bool +} + +type hostedGatewayTestTokenProvider func(context.Context) (string, error) + +func (f hostedGatewayTestTokenProvider) AccessToken(ctx context.Context) (string, error) { + return f(ctx) +} + +type hostedGatewayTestRoundTripper func(*http.Request) (*http.Response, error) + +func (f hostedGatewayTestRoundTripper) RoundTrip(r *http.Request) (*http.Response, error) { + return f(r) +} + +type hostedGatewayBootstrapWrite struct { + ledger hostedGatewayLedger + persisted bool + containsToken bool +} + +// The observer runs before the underlying socket writes any bootstrap bytes. +// Only the forward channel is wrapped, and the reverse acknowledgment arms it. +type hostedGatewayTestWriteObserver struct { + net.Conn + armed *atomic.Bool + before func() +} + +func (c *hostedGatewayTestWriteObserver) Write(data []byte) (int, error) { + if c.armed.CompareAndSwap(true, false) { + c.before() + } + return c.Conn.Write(data) +} + +type hostedGatewayTestFixture struct { + t *testing.T + settings hostedGatewaySettings + gateway *hostedGateway + options hostedGatewayTestOptions + server *httptest.Server + claims map[string]any + token atomic.Value + tokenCalls atomic.Int64 + httpCalls atomic.Int64 + creates atomic.Int64 + dials atomic.Int64 + bootstraps atomic.Int64 + callbacks atomic.Int64 + writeArmed atomic.Bool + writes chan hostedGatewayBootstrapWrite + reverse chan *http2.ClientConn + mu sync.Mutex + exists bool + requests []string + hellos []hostedHello + peers []*websocket.Conn + channelLogs *hostedTestChannelLog +} + +func newHostedGatewayTestFixture(t *testing.T, options hostedGatewayTestOptions) *hostedGatewayTestFixture { + t.Helper() + p := newHostedProtocolFixture(t) + f := &hostedGatewayTestFixture{ + t: t, options: options, exists: options.preexistingSession, + writes: make(chan hostedGatewayBootstrapWrite, 1), reverse: make(chan *http2.ClientConn, 1), + claims: map[string]any{"aud": "https://ai.azure.com", "tid": uuid.NewString(), + "oid": uuid.NewString(), "appid": uuid.NewString()}, + } + f.token.Store(hostedGatewayTestJWT(f.claims)) + f.server = httptest.NewServer(http.HandlerFunc(f.serveHTTP)) + local, err := url.Parse(f.server.URL) + if err != nil { + t.Fatal("could not prepare local gateway fixture") + } + f.settings = hostedGatewaySettings{ + config: hostedGatewayConfig{ + Protocol: hostedProtocol, Image: p.config, + ContainerImage: "example.invalid/hosted@" + brokerSHA([]byte("gateway fixture image")), + SessionID: p.hello.Challenge.SessionID, RuntimeProfileDigest: brokerSHA([]byte("gateway fixture profile")), + RuntimeEnvironment: maps.Clone(p.bootstrap.Environment), + OrkaBaseURL: f.server.URL, BrokerBaseURL: f.server.URL, + }, + bootstrap: p.bootstrap, signingKey: bytes.Clone(p.key), stateDir: filepath.Join(t.TempDir(), "gateway"), + } + if validateHostedGatewayConfig(f.settings.config) != nil { + t.Fatal("gateway fixture configuration is invalid") + } + store, ledger, err := openHostedGatewayStore(f.settings.stateDir, f.settings.config) + if err != nil { + t.Fatal("could not open gateway fixture store") + } + ctx, cancel := context.WithCancel(context.Background()) + provider := hostedGatewayTestTokenProvider(func(ctx context.Context) (string, error) { + f.tokenCalls.Add(1) + if ctx.Err() != nil { + return "", ctx.Err() + } + return f.token.Load().(string), nil + }) + transport := newHostedLocalTransport() + remote, _ := url.Parse(f.settings.config.Image.Target.ProjectEndpoint) + client := &http.Client{Timeout: 5 * time.Second, + CheckRedirect: func(*http.Request, []*http.Request) error { return errHostedInvalid }, + Transport: hostedGatewayTestRoundTripper(func(r *http.Request) (*http.Response, error) { + if r.URL.Scheme != "https" || r.URL.Host != remote.Host || r.GetBody != nil { + t.Error("gateway changed the configured authority or exposed a replayable request body") + return nil, errHostedInvalid + } + copy := r.Clone(r.Context()) + copy.URL.Scheme, copy.URL.Host, copy.Host = local.Scheme, local.Host, remote.Host + return transport.RoundTrip(copy) + })} + f.gateway = &hostedGateway{cfg: f.settings.config, provider: provider, httpClient: client, + store: store, ledger: ledger, ctx: ctx, cancel: cancel} + f.channelLogs = newHostedTestChannelLog() + f.gateway.channelLog = f.channelLogs.logger + f.gateway.dial = func(ctx context.Context, endpoint string, headers http.Header) (*websocket.Conn, *http.Response, error) { + index := f.dials.Add(1) + expected := "wss://" + remote.Host + remote.Path + "/agents/" + p.config.Target.AgentName + + "/endpoint/protocols/invocations_ws?api-version=v1&agent_session_id=" + f.settings.config.SessionID + if endpoint != expected || headers.Get("Authorization") != "Bearer "+f.token.Load().(string) || + headers.Get("Foundry-Features") != "HostedAgents=V1Preview" { + t.Error("gateway WebSocket route, exact session or Azure authorization changed") + return nil, nil, errHostedInvalid + } + dialer := websocket.Dialer{HandshakeTimeout: 5 * time.Second, + NetDialContext: func(ctx context.Context, network, address string) (net.Conn, error) { + if address != local.Host { + return nil, errHostedInvalid + } + conn, err := (&net.Dialer{Timeout: 5 * time.Second}).DialContext(ctx, network, address) + if err != nil || index != 1 { + return conn, err + } + return &hostedGatewayTestWriteObserver{Conn: conn, armed: &f.writeArmed, before: f.observeBootstrapWrite}, nil + }} + return dialer.DialContext(ctx, "ws://"+local.Host+strings.TrimPrefix(endpoint, "wss://"+remote.Host), headers) + } + t.Cleanup(func() { + f.gateway.close() + f.mu.Lock() + peers := append([]*websocket.Conn(nil), f.peers...) + f.mu.Unlock() + for _, peer := range peers { + _ = peer.Close() + } + transport.CloseIdleConnections() + f.server.Close() + clear(f.settings.signingKey) + }) + return f +} + +func hostedGatewayTestJWT(claims map[string]any) string { + data, _ := json.Marshal(claims) + return "fixture-header." + base64.RawURLEncoding.EncodeToString(data) + ".fixture-signature" +} + +func (f *hostedGatewayTestFixture) serveHTTP(w http.ResponseWriter, r *http.Request) { + if r.URL.Path == brokerStatusPath || r.URL.Path == "/internal/v2/acp/mcp/tools/call" { + f.callbacks.Add(1) + if r.Header.Get("Authorization") != "Bearer fixture-operation-authorization" { + f.t.Error("reverse channel changed caller authorization") + } + if f.options.callback != nil { + f.options.callback.ServeHTTP(w, r) + return + } + w.WriteHeader(http.StatusNoContent) + return + } + target := f.settings.config.Image.Target + base := "/api/projects/test-project/agents/" + target.AgentName + if r.URL.Path == base+"/endpoint/protocols/invocations_ws" { + f.serveWebSocket(w, r) + return + } + f.httpCalls.Add(1) + f.mu.Lock() + f.requests = append(f.requests, r.Method+" "+r.URL.RequestURI()) + exists := f.exists + f.mu.Unlock() + if r.Header.Get("Authorization") != "Bearer "+f.token.Load().(string) || + r.Header.Get("Foundry-Features") != "HostedAgents=V1Preview" || r.URL.RawQuery != "api-version=v1" { + f.t.Error("gateway HTTP routing or Azure authorization changed") + http.Error(w, "fixture request rejected", http.StatusBadRequest) + return + } + stage, status := "", http.StatusOK + var body map[string]any + switch { + case r.Method == http.MethodGet && r.URL.Path == base: + stage = "agent" + body = map[string]any{"name": target.AgentName, "agent_endpoint": map[string]any{ + "authorization_schemes": []any{map[string]any{"type": "entra"}}}} + case r.Method == http.MethodGet && r.URL.Path == base+"/versions/"+target.AgentVersion: + stage = "version" + body = map[string]any{"name": target.AgentName, "version": target.AgentVersion, "status": "active", + "definition": map[string]any{"kind": "hosted", + "container_configuration": map[string]any{"image": f.settings.config.ContainerImage}, + "protocol_versions": []any{map[string]any{"protocol": "invocations_ws", "version": "2.0.0"}}}} + case r.Method == http.MethodGet && r.URL.Path == base+brokerSessionSuffix(f.settings.config.SessionID): + stage, body = "session-get", f.sessionBody() + if !exists { + status, body = http.StatusNotFound, map[string]any{} + } + case r.Method == http.MethodPost && r.URL.Path == base+"/endpoint/sessions": + f.creates.Add(1) + data, err := io.ReadAll(io.LimitReader(r.Body, hostedMaxHandshakeBytes+1)) + var request brokerRemoteSession + if err != nil || acpDecode(data, &request, true) != nil || + request.ID != f.settings.config.SessionID || request.Version.Type != "version_ref" || request.Version.Version != target.AgentVersion { + f.t.Error("session create omitted the exact session or concrete version") + http.Error(w, "fixture create rejected", http.StatusBadRequest) + return + } + f.mu.Lock() + f.exists = true + f.mu.Unlock() + if f.options.ambiguousCreate { + conn, _, err := w.(http.Hijacker).Hijack() + if err == nil { + _ = conn.Close() + } + return + } + stage, status, body = "session-create", http.StatusCreated, f.sessionBody() + default: + f.t.Error("gateway called an unexpected Foundry API route") + http.Error(w, "fixture route rejected", http.StatusNotFound) + return + } + if f.options.remote != nil && status != http.StatusNotFound { + f.options.remote(stage, body) + } + w.Header().Set("Content-Type", "application/json") + w.WriteHeader(status) + _ = json.NewEncoder(w).Encode(body) +} + +func (f *hostedGatewayTestFixture) sessionBody() map[string]any { + return map[string]any{"agent_session_id": f.settings.config.SessionID, "status": "active", + "version_indicator": map[string]any{"type": "version_ref", "agent_version": f.settings.config.Image.Target.AgentVersion}} +} + +func (f *hostedGatewayTestFixture) serveWebSocket(w http.ResponseWriter, r *http.Request) { + upgrader := websocket.Upgrader{} + ws, err := upgrader.Upgrade(w, r, nil) + if err != nil { + return + } + defer ws.Close() + f.mu.Lock() + f.peers = append(f.peers, ws) + index := len(f.peers) + f.mu.Unlock() + _ = ws.SetReadDeadline(time.Now().Add(5 * time.Second)) + _ = ws.SetWriteDeadline(time.Now().Add(5 * time.Second)) + challenge := hostedChallenge{ + Protocol: hostedProtocol, DeploymentID: f.settings.config.Image.DeploymentID, + ConfigurationDigest: brokerJSONDigest(f.settings.config.Image), + AgentName: f.settings.config.Image.Target.AgentName, AgentVersion: f.settings.config.Image.Target.AgentVersion, + SessionID: f.settings.config.SessionID, BootID: f.settings.config.Image.DeploymentID, + Nonce: base64.RawURLEncoding.EncodeToString(bytes.Repeat([]byte{11}, 32)), + } + if f.options.challenge != nil { + f.options.challenge(index, &challenge) + } + if ws.WriteJSON(challenge) != nil { + return + } + var hello hostedHello + if readHostedWSJSON(ws, &hello) != nil { + return + } + if verifyHostedHello(hello, challenge, f.settings.config.Image.SigningPublicKey, time.Now()) != nil { + f.t.Error("gateway sent an invalid signed handshake") + return + } + f.mu.Lock() + f.hellos = append(f.hellos, hello) + f.mu.Unlock() + role := "forward" + if index == 2 { + role = "reverse" + } + if hello.Role != role { + f.t.Error("gateway did not assign distinct forward and reverse roles") + return + } + ack := hostedAccepted{Protocol: hostedProtocol, PairID: hello.PairID, Role: role, + BootID: challenge.BootID, BootstrapDigest: hello.BootstrapDigest} + if f.options.accepted != nil { + f.options.accepted(&ack) + } + if role == "reverse" { + f.writeArmed.Store(true) + if f.options.beforeReverseAck != nil { + f.options.beforeReverseAck() + } + } + if ws.WriteJSON(ack) != nil { + return + } + if role == "reverse" { + conn := newHostedWSConn(ws) + client, err := newHostedHTTP2ClientConn(conn) + if err != nil { + return + } + defer client.Close() + f.reverse <- client + select { + case <-f.gateway.ctx.Done(): + case <-conn.Done(): + } + return + } + kind, data, err := ws.ReadMessage() + if err != nil { + return + } + f.bootstraps.Add(1) + var bootstrap hostedBootstrap + valid := kind == websocket.TextMessage && brokerSHA(data) == hello.BootstrapDigest && + acpDecode(data, &bootstrap, true) == nil && reflect.DeepEqual(bootstrap, f.settings.bootstrap) + clear(data) + if !valid { + f.t.Error("gateway changed the signed bootstrap body") + return + } + if f.options.dropAfterBootstrap { + return + } + ready := hostedReady{Protocol: hostedProtocol, PairID: hello.PairID, BootID: challenge.BootID, Ready: true} + if f.options.ready != nil { + f.options.ready(&ready) + } + if ws.WriteJSON(ready) != nil { + return + } + serveHostedHTTP2(f.gateway.ctx, newHostedWSConn(ws), http.HandlerFunc(f.serveSupervisor)) +} + +func (f *hostedGatewayTestFixture) serveSupervisor(w http.ResponseWriter, r *http.Request) { + if r.URL.Path == "/v2/capabilities" { + value := map[string]any{"protocol": "orka.harness.v2", "transport": "http+ndjson", + "runtimeProfileDigest": f.settings.config.RuntimeProfileDigest, + "adapterDigests": map[string]string{"foundry-serve-acp": f.settings.config.RuntimeEnvironment["ORKA_ACP_FOUNDRY_ADAPTER_DIGEST"]}} + if f.options.capabilities != nil { + f.options.capabilities(value) + } + w.Header().Set("Content-Type", "application/json") + _ = json.NewEncoder(w).Encode(value) + return + } + if f.options.supervisor != nil { + f.options.supervisor.ServeHTTP(w, r) + return + } + w.WriteHeader(http.StatusNoContent) +} + +func (f *hostedGatewayTestFixture) observeBootstrapWrite() { + data, err := os.ReadFile(filepath.Join(f.settings.stateDir, "state.json")) + var observed hostedGatewayBootstrapWrite + observed.persisted = err == nil && acpDecode(data, &observed.ledger, true) == nil && + hostedGatewayLedgerValid(observed.ledger, f.settings.config) + for _, token := range []string{f.settings.bootstrap.ControllerToken, f.settings.bootstrap.CapabilitySecret, f.settings.bootstrap.ProviderToken} { + observed.containsToken = observed.containsToken || bytes.Contains(data, []byte(token)) + } + f.writes <- observed +} + +func (f *hostedGatewayTestFixture) initialize() error { + err := f.gateway.initialize(f.settings) + if err != nil { + f.gateway.close() + } + return err +} + +func (f *hostedGatewayTestFixture) assertNoBootstrap(t *testing.T) { + t.Helper() + if f.bootstraps.Load() != 0 { + t.Fatal("gateway disclosed bootstrap after failed admission") + } + select { + case <-f.writes: + t.Fatal("gateway wrote bootstrap bytes after failed admission") + default: + } +} + +func TestHostedGatewayChecksEncodedBootstrapBeforeRemoteIO(t *testing.T) { + for _, test := range []struct { + name string + field int + fits bool + }{ + {name: "escaped controller", field: 0}, + {name: "escaped capability", field: 1}, + {name: "escaped provider", field: 2}, + {name: "maximum raw credentials without expansion", fits: true}, + } { + t.Run(test.name, func(t *testing.T) { + f := newHostedGatewayTestFixture(t, hostedGatewayTestOptions{}) + fields := []*string{&f.settings.bootstrap.ControllerToken, &f.settings.bootstrap.CapabilitySecret, &f.settings.bootstrap.ProviderToken} + if test.fits { + for _, field := range fields { + *field = strings.Repeat("x", 16<<10) + } + } else { + *fields[test.field] = strings.Repeat("&", 16<<10) + } + if validateHostedBootstrap(f.settings.config.Image, f.settings.bootstrap) != nil { + t.Fatal("fixture did not supply valid raw bootstrap fields") + } + body, err := json.Marshal(f.settings.bootstrap) + if err != nil || (len(body) <= hostedMaxHandshakeBytes) != test.fits { + t.Fatal("fixture did not exercise the encoded handshake bound") + } + clear(body) + path := filepath.Join(f.settings.stateDir, "state.json") + before, err := os.ReadFile(path) + if err != nil { + t.Fatal("could not capture the initial ownership ledger") + } + err = f.initialize() + if test.fits { + if err != nil || !f.gateway.ready.Load() || f.creates.Load() != 1 || f.dials.Load() != 2 || f.bootstraps.Load() != 1 { + t.Fatal("valid encoded bootstrap did not establish exactly one lifetime") + } + return + } + if !errors.Is(err, errHostedInvalid) { + t.Fatal("oversized encoded bootstrap was not rejected") + } + if f.tokenCalls.Load() != 0 || f.httpCalls.Load() != 0 || f.creates.Load() != 0 || f.dials.Load() != 0 { + t.Error("oversized encoded bootstrap reached remote I/O") + } + after, err := os.ReadFile(path) + if err != nil || !bytes.Equal(before, after) { + t.Error("oversized encoded bootstrap changed durable ownership") + } + f.assertNoBootstrap(t) + }) + } +} + +func TestHostedGatewayBindsLifetimeAndPersistsBeforeBootstrap(t *testing.T) { + f := newHostedGatewayTestFixture(t, hostedGatewayTestOptions{}) + if f.initialize() != nil || !f.gateway.ready.Load() { + t.Fatal("valid local hosted lifetime did not become ready") + } + var observed hostedGatewayBootstrapWrite + select { + case observed = <-f.writes: + default: + t.Fatal("bootstrap write was not observed") + } + if !observed.persisted || observed.containsToken || !observed.ledger.ExposurePossible || + observed.ledger.Ready || observed.ledger.Closed || !observed.ledger.CreateAttempted || !observed.ledger.SessionCreated || + observed.ledger.PrincipalDigest == "" || observed.ledger.BootstrapDigest != brokerJSONDigest(f.settings.bootstrap) { + t.Fatal("complete private exposure record was not durable before the first bootstrap write") + } + f.mu.Lock() + hellos, requests := append([]hostedHello(nil), f.hellos...), append([]string(nil), f.requests...) + f.mu.Unlock() + if len(hellos) != 2 || hellos[0].Role != "forward" || hellos[1].Role != "reverse" || + hellos[0].PairID != hellos[1].PairID || hellos[0].Challenge != hellos[1].Challenge || + hellos[0].BootstrapDigest != hellos[1].BootstrapDigest || observed.ledger.Challenge != hellos[0].Challenge || + observed.ledger.PairID != hellos[0].PairID || f.dials.Load() != 2 || f.bootstraps.Load() != 1 { + t.Fatal("forward and reverse channels did not bind one exact lifetime and bootstrap") + } + base := "/api/projects/test-project/agents/" + f.settings.config.Image.Target.AgentName + want := []string{"GET " + base + "?api-version=v1", + "GET " + base + "/versions/" + f.settings.config.Image.Target.AgentVersion + "?api-version=v1", + "GET " + base + brokerSessionSuffix(f.settings.config.SessionID) + "?api-version=v1", + "POST " + base + "/endpoint/sessions?api-version=v1", + "GET " + base + brokerSessionSuffix(f.settings.config.SessionID) + "?api-version=v1"} + if !reflect.DeepEqual(requests, want) || f.creates.Load() != 1 { + t.Fatal("gateway did not validate, create and confirm exactly the configured session") + } + front := httptest.NewServer(f.gateway) + defer front.Close() + client := &http.Client{Timeout: 5 * time.Second} + for path, status := range map[string]int{"/healthz": http.StatusOK, "/v2/health": http.StatusNoContent} { + response, err := client.Get(front.URL + path) + if err != nil { + t.Fatal("ready gateway did not serve its HTTP interface") + } + _ = response.Body.Close() + if response.StatusCode != status { + t.Fatal("gateway changed the supervisor HTTP status") + } + } + var reverse *http2.ClientConn + select { + case reverse = <-f.reverse: + case <-time.After(5 * time.Second): + t.Fatal("reverse HTTP/2 channel did not connect") + } + for _, endpoint := range []struct{ method, url string }{ + {http.MethodGet, "http://broker" + brokerStatusPath}, + {http.MethodPost, "http://orka/internal/v2/acp/mcp/tools/call"}, + } { + ctx, cancel := context.WithTimeout(context.Background(), 5*time.Second) + request, _ := http.NewRequestWithContext(ctx, endpoint.method, endpoint.url, nil) + request.Header.Set("Authorization", "Bearer fixture-operation-authorization") + response, err := reverse.RoundTrip(request) + if err != nil { + cancel() + t.Fatal("reverse callback did not reach the fixed upstream") + } + _ = response.Body.Close() + cancel() + if response.StatusCode != http.StatusNoContent { + t.Fatal("reverse callback changed the upstream status") + } + } + if f.callbacks.Load() != 2 { + t.Fatal("reverse callback was lost or replayed") + } +} + +func TestHostedGatewayRejectsRemoteTargetDriftBeforeChannels(t *testing.T) { + for _, test := range []struct { + name, stage string + mutate func(map[string]any) + }{ + {"agent name", "agent", func(v map[string]any) { v["name"] = "other-agent" }}, + {"anonymous endpoint", "agent", func(v map[string]any) { v["agent_endpoint"] = map[string]any{"authorization_schemes": []any{}} }}, + {"non-Entra endpoint", "agent", func(v map[string]any) { + v["agent_endpoint"] = map[string]any{"authorization_schemes": []any{map[string]any{"type": "key"}}} + }}, + {"version name", "version", func(v map[string]any) { v["name"] = "other-agent" }}, + {"version number", "version", func(v map[string]any) { v["version"] = "9" }}, + {"inactive version", "version", func(v map[string]any) { v["status"] = "inactive" }}, + {"definition kind", "version", func(v map[string]any) { v["definition"].(map[string]any)["kind"] = "prompt" }}, + {"container digest", "version", func(v map[string]any) { + v["definition"].(map[string]any)["container_configuration"] = map[string]any{"image": "example.invalid/hosted@" + brokerSHA([]byte("other image"))} + }}, + {"missing WebSocket", "version", func(v map[string]any) { v["definition"].(map[string]any)["protocol_versions"] = []any{} }}, + {"WebSocket protocol version", "version", func(v map[string]any) { + v["definition"].(map[string]any)["protocol_versions"] = []any{map[string]any{"protocol": "invocations_ws", "version": "1.0.0"}} + }}, + {"duplicate WebSocket protocol", "version", func(v map[string]any) { + v["definition"].(map[string]any)["protocol_versions"] = []any{map[string]any{"protocol": "invocations_ws", "version": "2.0.0"}, map[string]any{"protocol": "invocations_ws", "version": "2.0.0"}} + }}, + {"created session ID", "session-create", func(v map[string]any) { v["agent_session_id"] = uuid.NewString() }}, + {"created session version", "session-create", func(v map[string]any) { v["version_indicator"].(map[string]any)["agent_version"] = "9" }}, + {"session version selector", "session-get", func(v map[string]any) { v["version_indicator"].(map[string]any)["type"] = "latest" }}, + {"confirmed session ID", "session-get", func(v map[string]any) { v["agent_session_id"] = uuid.NewString() }}, + {"inactive session", "session-get", func(v map[string]any) { v["status"] = "inactive" }}, + } { + t.Run(test.name, func(t *testing.T) { + f := newHostedGatewayTestFixture(t, hostedGatewayTestOptions{remote: func(stage string, body map[string]any) { + if stage == test.stage { + test.mutate(body) + } + }}) + if f.initialize() == nil || f.gateway.ready.Load() || f.dials.Load() != 0 { + t.Fatal("changed remote target reached hosted channel admission") + } + f.assertNoBootstrap(t) + }) + } + t.Run("preexisting session", func(t *testing.T) { + f := newHostedGatewayTestFixture(t, hostedGatewayTestOptions{preexistingSession: true}) + if f.initialize() == nil || f.creates.Load() != 0 || f.dials.Load() != 0 { + t.Fatal("gateway adopted or replaced an existing session without creation ownership") + } + f.assertNoBootstrap(t) + }) +} + +func TestHostedGatewayRejectsChallengeAndChannelBindingDrift(t *testing.T) { + for name, mutate := range map[string]func(*hostedChallenge){ + "protocol": func(c *hostedChallenge) { c.Protocol = "other" }, + "deployment": func(c *hostedChallenge) { c.DeploymentID = uuid.NewString() }, + "configuration": func(c *hostedChallenge) { c.ConfigurationDigest = brokerSHA([]byte("other configuration")) }, + "agent": func(c *hostedChallenge) { c.AgentName = "other-agent" }, + "version": func(c *hostedChallenge) { c.AgentVersion = "9" }, + "session": func(c *hostedChallenge) { c.SessionID = uuid.NewString() }, + "invalid boot": func(c *hostedChallenge) { c.BootID = uuid.Nil.String() }, + "invalid nonce": func(c *hostedChallenge) { c.Nonce += "=" }, + } { + t.Run(name, func(t *testing.T) { + f := newHostedGatewayTestFixture(t, hostedGatewayTestOptions{challenge: func(_ int, c *hostedChallenge) { mutate(c) }}) + if f.initialize() == nil || f.dials.Load() != 1 { + t.Fatal("gateway admitted a challenge outside the exact configured target") + } + f.assertNoBootstrap(t) + }) + } + for name, mutate := range map[string]func(*hostedChallenge){ + "cross-boot pair": func(c *hostedChallenge) { c.BootID = uuid.NewString() }, + "cross-nonce pair": func(c *hostedChallenge) { c.Nonce = base64.RawURLEncoding.EncodeToString(bytes.Repeat([]byte{12}, 32)) }, + } { + t.Run(name, func(t *testing.T) { + f := newHostedGatewayTestFixture(t, hostedGatewayTestOptions{challenge: func(index int, c *hostedChallenge) { + if index == 2 { + mutate(c) + } + }}) + if f.initialize() == nil || f.dials.Load() != 2 { + t.Fatal("gateway paired channels from different hosted challenges") + } + f.mu.Lock() + hellos := len(f.hellos) + f.mu.Unlock() + if hellos != 0 { + t.Fatal("gateway signed a role before matching both hosted challenges") + } + f.assertNoBootstrap(t) + }) + } + for name, mutate := range map[string]func(*hostedAccepted){ + "protocol": func(a *hostedAccepted) { a.Protocol = "other" }, + "pair": func(a *hostedAccepted) { a.PairID = uuid.NewString() }, + "role": func(a *hostedAccepted) { a.Role = "other" }, + "boot": func(a *hostedAccepted) { a.BootID = uuid.NewString() }, + "bootstrap": func(a *hostedAccepted) { a.BootstrapDigest = brokerSHA([]byte("other bootstrap")) }, + } { + for _, role := range []string{"forward", "reverse"} { + t.Run(role+" acknowledgment "+name, func(t *testing.T) { + f := newHostedGatewayTestFixture(t, hostedGatewayTestOptions{accepted: func(a *hostedAccepted) { + if a.Role == role { + mutate(a) + } + }}) + if f.initialize() == nil { + t.Fatal("gateway accepted a changed signed-role acknowledgment") + } + f.assertNoBootstrap(t) + }) + } + } +} + +func TestHostedGatewayAmbiguousCreateIsNeverReplayedOrAdopted(t *testing.T) { + f := newHostedGatewayTestFixture(t, hostedGatewayTestOptions{ambiguousCreate: true}) + if f.initialize() == nil || f.creates.Load() != 1 || f.dials.Load() != 0 { + t.Fatal("gateway continued after losing the create response") + } + f.assertNoBootstrap(t) + store, ledger, err := openHostedGatewayStore(f.settings.stateDir, f.settings.config) + if err != nil || !ledger.CreateAttempted || ledger.SessionCreated || ledger.ExposurePossible { + t.Fatal("ambiguous creation intent did not survive restart") + } + defer store.close() + httpBefore, tokensBefore := f.httpCalls.Load(), f.tokenCalls.Load() + restarted := &hostedGateway{cfg: f.settings.config, provider: f.gateway.provider, + httpClient: f.gateway.httpClient, store: store, ledger: ledger} + for range 2 { + if restarted.ensureSession(context.Background()) == nil { + t.Fatal("gateway adopted a remotely active session after an ambiguous create") + } + } + if f.httpCalls.Load() != httpBefore || f.tokenCalls.Load() != tokensBefore || f.creates.Load() != 1 { + t.Fatal("ambiguous create caused a token request, discovery request or replay") + } +} + +func TestHostedGatewayPersistenceFailurePreventsBootstrapWrite(t *testing.T) { + f := newHostedGatewayTestFixture(t, hostedGatewayTestOptions{}) + f.options.beforeReverseAck = func() { + if os.Rename(f.settings.stateDir, f.settings.stateDir+"-unavailable") != nil { + t.Error("could not inject exposure persistence failure") + } + } + if f.initialize() == nil || f.gateway.ready.Load() { + t.Fatal("gateway became ready after failing to persist exposure") + } + f.assertNoBootstrap(t) +} + +func TestHostedGatewayPrincipalDriftIsRejectedBeforeNetwork(t *testing.T) { + for _, claim := range []string{"aud", "tid", "oid", "appid", "azp"} { + t.Run(claim, func(t *testing.T) { + f := newHostedGatewayTestFixture(t, hostedGatewayTestOptions{}) + if f.gateway.validateRemote(context.Background()) != nil { + t.Fatal("valid initial principal was not enrolled") + } + changed := maps.Clone(f.claims) + changed[claim] = uuid.NewString() + if claim == "aud" { + changed[claim] = "https://other.invalid" + } + f.token.Store(hostedGatewayTestJWT(changed)) + before := f.httpCalls.Load() + if f.gateway.validateRemote(context.Background()) == nil || f.httpCalls.Load() != before || f.dials.Load() != 0 { + t.Fatal("a changed Azure principal reached remote HTTP or WebSocket calls") + } + }) + } + t.Run("same principal token refresh", func(t *testing.T) { + f := newHostedGatewayTestFixture(t, hostedGatewayTestOptions{}) + if f.gateway.validateRemote(context.Background()) != nil { + t.Fatal("valid initial principal was not enrolled") + } + refreshed := maps.Clone(f.claims) + refreshed["iat"], refreshed["exp"] = time.Now().Unix(), time.Now().Add(time.Hour).Unix() + f.token.Store(hostedGatewayTestJWT(refreshed)) + if f.gateway.validateRemote(context.Background()) != nil || f.httpCalls.Load() != 4 { + t.Fatal("normal token refresh changed the enrolled principal") + } + }) +} + +func TestHostedGatewayExposureBlocksRestartAndCredentialChangesBeforeNetwork(t *testing.T) { + f := newHostedGatewayTestFixture(t, hostedGatewayTestOptions{dropAfterBootstrap: true}) + if f.initialize() == nil || f.bootstraps.Load() != 1 { + t.Fatal("fixture did not lose transport after bootstrap exposure") + } + for name, mutate := range map[string]func(*hostedGatewaySettings){ + "unchanged restart": func(*hostedGatewaySettings) {}, + "controller credential": func(s *hostedGatewaySettings) { s.bootstrap.ControllerToken = strings.Repeat("changed-controller-", 3) }, + "capability credential": func(s *hostedGatewaySettings) { + s.bootstrap.CapabilitySecret = strings.Repeat("changed-capability-", 3) + }, + "provider credential": func(s *hostedGatewaySettings) { s.bootstrap.ProviderToken = strings.Repeat("changed-provider-", 3) }, + "configuration": func(s *hostedGatewaySettings) { s.config.RuntimeEnvironment["ORKA_ACP_CONTROLLER_EPOCH"] = "2" }, + } { + t.Run(name, func(t *testing.T) { + settings := f.settings + settings.signingKey = bytes.Clone(f.settings.signingKey) + settings.config.RuntimeEnvironment = maps.Clone(f.settings.config.RuntimeEnvironment) + settings.bootstrap.Environment = maps.Clone(f.settings.bootstrap.Environment) + mutate(&settings) + var called atomic.Int64 + provider := hostedGatewayTestTokenProvider(func(context.Context) (string, error) { + called.Add(1) + return "", errors.New("fixture forbids external network") + }) + gateway, err := newHostedGateway(context.Background(), settings, provider) + if gateway != nil { + gateway.close() + } + if err == nil || called.Load() != 0 { + t.Fatal("exposed ownership or credential drift reached Azure token acquisition") + } + }) + } +} + +func TestHostedGatewayRejectsReadyAndCapabilityDriftAfterExposure(t *testing.T) { + for name, mutate := range map[string]func(*hostedReady){ + "protocol": func(r *hostedReady) { r.Protocol = "other" }, + "pair": func(r *hostedReady) { r.PairID = uuid.NewString() }, + "boot": func(r *hostedReady) { r.BootID = uuid.NewString() }, + "not ready": func(r *hostedReady) { r.Ready = false }, + } { + t.Run("ready "+name, func(t *testing.T) { + f := newHostedGatewayTestFixture(t, hostedGatewayTestOptions{ready: mutate}) + if f.initialize() == nil || f.gateway.ready.Load() || f.bootstraps.Load() != 1 { + t.Fatal("gateway became ready after a changed hosted-ready acknowledgment") + } + }) + } + for name, mutate := range map[string]func(map[string]any){ + "protocol": func(v map[string]any) { v["protocol"] = "orka.harness.v1" }, + "transport": func(v map[string]any) { v["transport"] = "other" }, + "profile": func(v map[string]any) { v["runtimeProfileDigest"] = brokerSHA([]byte("other profile")) }, + "adapter": func(v map[string]any) { + v["adapterDigests"] = map[string]string{"foundry-serve-acp": brokerSHA([]byte("other adapter"))} + }, + } { + t.Run("capabilities "+name, func(t *testing.T) { + f := newHostedGatewayTestFixture(t, hostedGatewayTestOptions{capabilities: mutate}) + if f.initialize() == nil || f.gateway.ready.Load() || f.bootstraps.Load() != 1 { + t.Fatal("gateway became ready with a different supervisor profile or adapter") + } + }) + } +} + +func TestHostedGatewayStreamsWhileOtherRequestsProgress(t *testing.T) { + release := make(chan struct{}) + var streams atomic.Int64 + f := newHostedGatewayTestFixture(t, hostedGatewayTestOptions{supervisor: http.HandlerFunc(func(w http.ResponseWriter, r *http.Request) { + if r.URL.Path != "/v2/events" { + w.WriteHeader(http.StatusNoContent) + return + } + streams.Add(1) + w.Header().Set("Content-Type", "application/x-ndjson") + w.Header().Set("Trailer", "X-Stream-Complete") + w.WriteHeader(http.StatusAccepted) + _, _ = io.WriteString(w, "{\"part\":1}\n") + w.(http.Flusher).Flush() + select { + case <-release: + case <-r.Context().Done(): + return + } + _, _ = io.WriteString(w, "{\"part\":2}\n") + w.Header().Set("X-Stream-Complete", "yes") + })}) + if f.initialize() != nil { + t.Fatal("streaming fixture gateway did not become ready") + } + front := httptest.NewServer(f.gateway) + defer front.Close() + client := &http.Client{Timeout: 5 * time.Second} + response, err := client.Get(front.URL + "/v2/events") + if err != nil { + t.Fatal("gateway buffered the stream before delivering response headers") + } + defer response.Body.Close() + reader := bufio.NewReader(response.Body) + first, err := reader.ReadString('\n') + if err != nil || first != "{\"part\":1}\n" || response.StatusCode != http.StatusAccepted { + t.Fatal("gateway changed or buffered the first stream record") + } + other, err := client.Get(front.URL + "/v2/health") + if err != nil { + t.Fatal("open stream blocked another HTTP/2 request") + } + _ = other.Body.Close() + if other.StatusCode != http.StatusNoContent { + t.Fatal("concurrent request did not reach the supervisor") + } + close(release) + rest, err := io.ReadAll(reader) + if err != nil || string(rest) != "{\"part\":2}\n" || response.Trailer.Get("X-Stream-Complete") != "yes" || streams.Load() != 1 { + t.Fatal("gateway lost stream completion, trailers or request ownership") + } +} + +func TestHostedGatewayLostMutationClosesLifetimeWithoutReplay(t *testing.T) { + var accepted atomic.Int64 + var f *hostedGatewayTestFixture + f = newHostedGatewayTestFixture(t, hostedGatewayTestOptions{supervisor: http.HandlerFunc(func(w http.ResponseWriter, r *http.Request) { + _, _ = io.Copy(io.Discard, r.Body) + accepted.Add(1) + f.mu.Lock() + forward := f.peers[0] + f.mu.Unlock() + _ = forward.Close() + })}) + if f.initialize() != nil { + t.Fatal("mutation fixture gateway did not become ready") + } + httpBefore, tokensBefore := f.httpCalls.Load(), f.tokenCalls.Load() + front := httptest.NewServer(f.gateway) + defer front.Close() + client := &http.Client{Timeout: 5 * time.Second} + response, err := client.Post(front.URL+"/v2/sessions/fixture/operations", "application/json", strings.NewReader(`{"operation":"fixture"}`)) + if err != nil { + t.Fatal("gateway did not report the failed mutation transport") + } + _ = response.Body.Close() + if response.StatusCode != http.StatusBadGateway || accepted.Load() != 1 { + t.Fatal("ambiguous mutation was replayed or reported as successful") + } + select { + case <-f.gateway.ctx.Done(): + case <-time.After(5 * time.Second): + t.Fatal("lost hosted channel did not close the gateway lifetime") + } + f.gateway.close() + denied := httptest.NewRecorder() + f.gateway.ServeHTTP(denied, httptest.NewRequest(http.MethodPost, "/v2/sessions/fixture/operations", strings.NewReader(`{}`))) + if denied.Code != http.StatusServiceUnavailable || accepted.Load() != 1 || f.httpCalls.Load() != httpBefore || + f.tokenCalls.Load() != tokensBefore || f.dials.Load() != 2 || f.bootstraps.Load() != 1 { + t.Fatal("closed gateway retried work, reconnected or admitted a replacement request") + } + var ledger hostedGatewayLedger + if readHostedConfig(filepath.Join(f.settings.stateDir, "state.json"), &ledger) != nil || !ledger.ExposurePossible || !ledger.Closed { + t.Fatal("lost channel erased durable possible exposure") + } +} + +func hostedGatewayTestClosedWithoutReplay(t *testing.T, f *hostedGatewayTestFixture, httpBefore, tokensBefore, callbacks int64) { + t.Helper() + hostedTestDone(t, f.gateway.ctx.Done()) + f.gateway.close() + denied := httptest.NewRecorder() + f.gateway.ServeHTTP(denied, httptest.NewRequest(http.MethodPost, "/v2/sessions/fixture/operations", strings.NewReader(`{}`))) + if denied.Code != http.StatusServiceUnavailable || f.callbacks.Load() != callbacks || f.httpCalls.Load() != httpBefore || + f.tokenCalls.Load() != tokensBefore || f.creates.Load() != 1 || f.dials.Load() != 2 || f.bootstraps.Load() != 1 { + t.Fatal("platform-closed gateway retried, reconnected or admitted a new operation") + } + var ledger hostedGatewayLedger + if readHostedConfig(filepath.Join(f.settings.stateDir, "state.json"), &ledger) != nil || !ledger.ExposurePossible || !ledger.Closed { + t.Fatal("platform closure lost durable possible exposure") + } + settings := f.settings + settings.signingKey = bytes.Clone(f.settings.signingKey) + var tokens atomic.Int64 + provider := hostedGatewayTestTokenProvider(func(context.Context) (string, error) { + tokens.Add(1) + return "", errHostedInvalid + }) + restarted, err := newHostedGateway(context.Background(), settings, provider) + if restarted != nil { + restarted.close() + } + if err == nil || restarted != nil || tokens.Load() != 0 { + t.Fatal("platform-closed lifetime could reseed or reach remote authentication after restart") + } +} + +func TestHostedGatewayIdleGoingAwayClosesLifetimeWithoutReseed(t *testing.T) { + for index, role := range []string{"forward", "reverse"} { + t.Run(role, func(t *testing.T) { + f := newHostedGatewayTestFixture(t, hostedGatewayTestOptions{}) + earliest := time.Now() + if f.initialize() != nil { + t.Fatal("idle gateway fixture did not become ready") + } + latest := time.Now() + httpBefore, tokensBefore := f.httpCalls.Load(), f.tokenCalls.Load() + f.mu.Lock() + peer := f.peers[index] + f.mu.Unlock() + hostedTestGoingAway(t, peer) + hostedGatewayTestClosedWithoutReplay(t, f, httpBefore, tokensBefore, 0) + hostedTestClosedChannelLogs(t, f.channelLogs, role, earliest, latest) + }) + } +} + +func TestHostedGatewayReverseGoingAwayAfterAcceptanceDoesNotReplay(t *testing.T) { + var accepted atomic.Int64 + var f *hostedGatewayTestFixture + f = newHostedGatewayTestFixture(t, hostedGatewayTestOptions{callback: http.HandlerFunc(func(w http.ResponseWriter, r *http.Request) { + if _, err := io.Copy(io.Discard, r.Body); err != nil { + t.Error("synthetic tool request body was not accepted") + return + } + accepted.Add(1) + f.mu.Lock() + peer := f.peers[1] + f.mu.Unlock() + if peer.WriteControl(websocket.CloseMessage, websocket.FormatCloseMessage(websocket.CloseGoingAway, hostedTestCloseText), time.Now().Add(time.Second)) != nil { + t.Error("could not close the reverse channel after synthetic tool acceptance") + return + } + // Never return a successful response before transport loss is observed. + select { + case <-r.Context().Done(): + case <-time.After(3 * time.Second): + t.Error("reverse transport loss did not cancel the accepted callback") + } + })}) + earliest := time.Now() + if f.initialize() != nil { + t.Fatal("reverse acceptance fixture did not become ready") + } + latest := time.Now() + httpBefore, tokensBefore := f.httpCalls.Load(), f.tokenCalls.Load() + var reverse *http2.ClientConn + select { + case reverse = <-f.reverse: + case <-time.After(3 * time.Second): + t.Fatal("reverse fixture was not established") + } + ctx, cancel := context.WithTimeout(context.Background(), 5*time.Second) + defer cancel() + for attempt := range 2 { + request, _ := http.NewRequestWithContext(ctx, http.MethodPost, "http://orka/internal/v2/acp/mcp/tools/call", strings.NewReader(`{"tool":"synthetic"}`)) + request.GetBody = nil + request.Header.Set("Authorization", "Bearer fixture-operation-authorization") + response, err := reverse.RoundTrip(request) + if response != nil { + _ = response.Body.Close() + } + if err == nil || response != nil || accepted.Load() != 1 { + t.Fatal("ambiguous reverse mutation succeeded or was replayed") + } + if attempt == 0 { + hostedGatewayTestClosedWithoutReplay(t, f, httpBefore, tokensBefore, 1) + } + } + if f.callbacks.Load() != 1 { + t.Fatal("a second reverse callback reached the backend after channel closure") + } + hostedTestClosedChannelLogs(t, f.channelLogs, "reverse", earliest, latest) +} diff --git a/hosted_handshake_deadline_test.go b/hosted_handshake_deadline_test.go new file mode 100644 index 0000000..e2a27c3 --- /dev/null +++ b/hosted_handshake_deadline_test.go @@ -0,0 +1,154 @@ +package main + +import ( + "context" + "net" + "net/http" + "net/http/httptest" + "strings" + "testing" + "testing/synctest" + "time" + + "github.com/google/uuid" + "github.com/gorilla/websocket" +) + +// Real Gorilla handshakes over net.Pipe exercise the production 10/30-second +// deadlines without an external network service. +func hostedHandshakePipeDial(t *testing.T, ctx context.Context, handler http.Handler, endpoint string, headers http.Header) (*websocket.Conn, *http.Response, error) { + t.Helper() + serverConn, clientConn := net.Pipe() + listener := &hostedPipeListener{connection: serverConn, closed: make(chan struct{})} + server := &http.Server{Handler: handler} + served := make(chan struct{}) + go func() { _ = server.Serve(listener); close(served) }() + t.Cleanup(func() { + _ = clientConn.Close() + _ = serverConn.Close() + _ = server.Close() + <-served + }) + dialer := &websocket.Dialer{HandshakeTimeout: time.Minute, + NetDialContext: func(context.Context, string, string) (net.Conn, error) { return clientConn, nil }} + return dialer.DialContext(ctx, strings.Replace(endpoint, "wss://", "ws://", 1), headers) +} + +func TestHostedGatewayHandshakeDeadlinesArePerPhase(t *testing.T) { + for _, test := range []struct { + name string + secondChallenge, firstAck, lastAck time.Duration + }{ + {name: "second challenge after original write deadline", secondChallenge: 11 * time.Second}, + {name: "bootstrap after slow reverse acknowledgment", lastAck: 11 * time.Second}, + {name: "acknowledgment after original read deadline", secondChallenge: 21 * time.Second, firstAck: 11 * time.Second}, + } { + t.Run(test.name, func(t *testing.T) { + // The gateway reaches HTTP/2, whose shared channel pool cannot cross + // synctest bubbles. Exercise these independent deadlines in parallel + // using real time; the server-only cases below remain on fake time. + t.Parallel() + var f *hostedGatewayTestFixture + f = newHostedGatewayTestFixture(t, hostedGatewayTestOptions{ + challenge: func(index int, _ *hostedChallenge) { + f.mu.Lock() + peer := f.peers[index-1] + f.mu.Unlock() + // The fixture's original five-second deadline must not mask + // the production per-phase limits being exercised here. + _ = peer.SetReadDeadline(time.Now().Add(time.Minute)) + _ = peer.SetWriteDeadline(time.Now().Add(time.Minute)) + if index == 2 { + time.Sleep(test.secondChallenge) + } + }, + accepted: func(ack *hostedAccepted) { + if ack.Role == "forward" { + time.Sleep(test.firstAck) + } + }, + beforeReverseAck: func() { time.Sleep(test.lastAck) }, + }) + f.gateway.cancel() + f.gateway.ctx, f.gateway.cancel = context.WithCancel(t.Context()) + defer f.gateway.close() + // Reuse the remote metadata and WebSocket fixtures in memory; + // no DNS, network service, or Azure is used. + f.gateway.httpClient.Transport = hostedGatewayTestRoundTripper(func(r *http.Request) (*http.Response, error) { + response := httptest.NewRecorder() + f.serveHTTP(response, r) + return response.Result(), nil + }) + f.gateway.dial = func(ctx context.Context, endpoint string, headers http.Header) (*websocket.Conn, *http.Response, error) { + f.dials.Add(1) + return hostedHandshakePipeDial(t, ctx, http.HandlerFunc(f.serveWebSocket), endpoint, headers) + } + start := time.Now() + if f.initialize() != nil { + t.Fatalf("valid delayed handshake rejected: elapsed=%s exposure=%t bootstrapDeliveries=%d", + time.Since(start), f.gateway.ledger.ExposurePossible, f.bootstraps.Load()) + } + if time.Since(start) < test.secondChallenge+test.firstAck+test.lastAck || !f.gateway.ready.Load() || + !f.gateway.ledger.ExposurePossible || !f.gateway.ledger.Ready || + f.creates.Load() != 1 || f.dials.Load() != 2 || f.bootstraps.Load() != 1 { + t.Fatal("slow handshake skipped its delay, retried, or failed to establish one lifetime") + } + }) + } +} + +func TestHostedServerHandshakeDeadlinesArePerPhase(t *testing.T) { + for _, test := range []struct { + name string + challengeRead, hello, bootstrap time.Duration + }{ + {name: "hello after original write deadline", hello: 11 * time.Second}, + {name: "hello after slow challenge write", challengeRead: 9 * time.Second, hello: 22 * time.Second}, + {name: "bootstrap after original read deadline", hello: 2 * time.Second, bootstrap: 29 * time.Second}, + } { + t.Run(test.name, func(t *testing.T) { + f := newHostedServerTestFixture(t, false) + synctest.Test(t, func(t *testing.T) { + f.server.cancel() + f.server.ctx, f.server.cancel = context.WithCancel(t.Context()) + defer f.server.close() //nolint:errcheck + ws, response, err := hostedHandshakePipeDial(t, t.Context(), f.server, "ws://fixture.invalid/invocations_ws", nil) + if response != nil && response.Body != nil { + _ = response.Body.Close() + } + if err != nil { + t.Fatal("could not establish in-memory hosted handshake") + } + defer ws.Close() //nolint:errcheck + _ = ws.SetReadDeadline(time.Now().Add(time.Minute)) + start := time.Now() + time.Sleep(test.challengeRead) + var challenge hostedChallenge + if readHostedWSJSON(ws, &challenge) != nil || challenge != f.server.challenge { + t.Fatal("valid delayed challenge was rejected") + } + time.Sleep(test.hello) + pairID, digest := uuid.NewString(), f.protocol.hello.BootstrapDigest + if ws.WriteJSON(f.hello(t, "forward", pairID, digest)) != nil { + t.Fatalf("valid hello could not be delivered after %s", time.Since(start)) + } + var ack hostedAccepted + if readHostedWSJSON(ws, &ack) != nil || ack != (hostedAccepted{Protocol: hostedProtocol, + PairID: pairID, Role: "forward", BootID: f.server.challenge.BootID, BootstrapDigest: digest}) { + t.Fatalf("valid hello was not acknowledged after %s", time.Since(start)) + } + if test.bootstrap != 0 { + time.Sleep(test.bootstrap) + f.sendBootstrap(t, ws) + synctest.Wait() + if !f.state(func(pair *hostedPair) bool { return pair != nil && pair.bootstrap != nil }) { + t.Fatal("bootstrap within its own read window was rejected") + } + } + if f.calls.Load() != 0 || f.server.ctx.Err() != nil { + t.Fatal("incomplete pair launched a supervisor or closed its lifetime") + } + }) + }) + } +} diff --git a/hosted_lifecycle_test.go b/hosted_lifecycle_test.go new file mode 100644 index 0000000..4429ff9 --- /dev/null +++ b/hosted_lifecycle_test.go @@ -0,0 +1,181 @@ +package main + +import ( + "bufio" + "context" + "errors" + "io" + "maps" + "net" + "net/http" + "sync" + "testing" + "testing/synctest" + "time" +) + +func TestHostedHTTPWaitsForSupervisorShutdown(t *testing.T) { + for _, role := range []string{"forward", "reverse"} { + t.Run(role, func(t *testing.T) { + hostedServerTestHealth(t, http.StatusOK) + f := newHostedServerTestFixture(t, true) + stopRequested, releaseStop := make(chan struct{}), make(chan struct{}) + release := sync.OnceFunc(func() { close(releaseStop) }) + t.Cleanup(release) + f.server.runner = func(ctx context.Context, environment map[string]string) (<-chan error, error) { + f.calls.Add(1) + f.captured <- maps.Clone(environment) + done := make(chan error, 1) + go func() { + <-ctx.Done() + close(stopRequested) + <-releaseStop + done <- nil + close(done) + }() + return done, nil + } + forward, reverse, _ := hostedServerTestRunningPair(t, f) + returned := make(chan error, 1) + go func() { returned <- serveHostedHTTP(f.server.ctx, "127.0.0.1:0", f.server) }() + channel := forward + if role == "reverse" { + channel = reverse + } + hostedTestGoingAway(t, channel.ws) + hostedTestDone(t, stopRequested) + select { + case <-returned: + t.Fatal("hosted entry point returned before the supervisor finished shutdown") + case <-time.After(25 * time.Millisecond): + } + release() + select { + case err := <-returned: + if err != nil { + t.Fatal("clean supervisor shutdown was rejected") + } + case <-time.After(3 * time.Second): + t.Fatal("hosted entry point did not join supervisor shutdown") + } + if f.calls.Load() != 1 { + t.Fatal("platform closure relaunched the joined supervisor") + } + }) + } +} +func TestHostedHTTPStreamsBodyBeyondThirtySeconds(t *testing.T) { + synctest.Test(t, func(t *testing.T) { + serverConn, clientConn := net.Pipe() + listener := &hostedPipeListener{connection: serverConn, closed: make(chan struct{})} + server := newHostedHTTPServer(http.HandlerFunc(func(w http.ResponseWriter, r *http.Request) { + body, err := io.ReadAll(io.LimitReader(r.Body, 4)) + if err != nil || string(body) != "abc" { + http.Error(w, "stream interrupted", http.StatusBadRequest) + return + } + w.WriteHeader(http.StatusCreated) + })) + connectionClosed := make(chan struct{}) + server.ConnState = func(_ net.Conn, state http.ConnState) { + if state == http.StateClosed { + close(connectionClosed) + } + } + served := make(chan struct{}) + go func() { _ = server.Serve(listener); close(served) }() + written := make(chan struct{}) + var writeErr error + defer func() { + _ = server.Close() + _ = clientConn.Close() + <-served + <-written + <-connectionClosed + }() + go func() { + _, err := io.WriteString(clientConn, "PUT /artifact HTTP/1.1\r\nHost: fixture\r\nContent-Length: 3\r\n\r\na") + if err == nil { + time.Sleep(31 * time.Second) + _, err = io.WriteString(clientConn, "bc") + } + writeErr = err + close(written) + }() + response, err := http.ReadResponse(bufio.NewReader(clientConn), &http.Request{Method: http.MethodPut}) + if err != nil { + t.Fatal("streaming request did not receive its response") + } + defer response.Body.Close() //nolint:errcheck + if response.StatusCode != http.StatusCreated { + t.Fatalf("valid slow body status = %d, want 201", response.StatusCode) + } + <-written + if writeErr != nil { + t.Fatal("valid slow body write failed") + } + }) +} + +func TestHostedSupervisorJoinReportsUnprovenShutdown(t *testing.T) { + for _, timeout := range []bool{false, true} { + name := "supervisor cleanup failed" + if timeout { + name = "supervisor did not finish" + } + t.Run(name, func(t *testing.T) { + synctest.Test(t, func(t *testing.T) { + ctx, cancel := context.WithCancel(t.Context()) + defer cancel() + server := &hostedServer{ctx: ctx, cancel: cancel} + pair := &hostedPair{server: server, running: true, done: make(chan struct{}), stopped: make(chan struct{})} + server.pair = pair + result := make(chan error, 1) + if !timeout { + result <- errHostedInvalid + } + started := time.Now() + go func() { pair.joinSupervisor(result); close(pair.stopped) }() + if err := server.close(); !errors.Is(err, errHostedInvalid) { + t.Fatal("unproven supervisor shutdown was accepted") + } + if timeout && time.Since(started) != hostedSupervisorShutdownWait { + t.Fatal("supervisor join did not use its bounded shutdown window") + } + }) + }) + } +} + +func TestHostedHTTPStreamingRetainsHeaderAndIdleBounds(t *testing.T) { + server := newHostedHTTPServer(http.NotFoundHandler()) + if server.ReadTimeout != 0 || server.WriteTimeout != 0 { + t.Fatal("hosted proxy applies a whole-stream timeout") + } + if server.ReadHeaderTimeout != 10*time.Second || server.IdleTimeout != 2*time.Minute || server.MaxHeaderBytes != 32<<10 { + t.Fatal("hosted proxy lost its header or idle bounds") + } +} + +type hostedPipeListener struct { + connection net.Conn + accepted bool + closed chan struct{} + once sync.Once +} + +func (l *hostedPipeListener) Accept() (net.Conn, error) { + if !l.accepted { + l.accepted = true + return l.connection, nil + } + <-l.closed + return nil, net.ErrClosed +} + +func (l *hostedPipeListener) Close() error { + l.once.Do(func() { close(l.closed) }) + return nil +} + +func (*hostedPipeListener) Addr() net.Addr { return &net.TCPAddr{} } diff --git a/hosted_main.go b/hosted_main.go new file mode 100644 index 0000000..e1ea745 --- /dev/null +++ b/hosted_main.go @@ -0,0 +1,126 @@ +package main + +import ( + "context" + "errors" + "flag" + "io" + "net" + "net/http" + "os" + "os/signal" + "strconv" + "syscall" + "time" +) + +func maybeServeHosted(args []string) (bool, error) { + mode := "" + for index, arg := range args { + for _, value := range []string{"hosted", "hosted-gateway"} { + if arg == "--protocol="+value || (arg == "--protocol" && index+1 < len(args) && args[index+1] == value) { + mode = value + } + } + } + if mode == "" { + return false, nil + } + flags := flag.NewFlagSet("foundry-hosted", flag.ContinueOnError) + flags.SetOutput(io.Discard) + protocol := flags.String("protocol", "", "") + defaultPath := hostedImageConfigPath + if mode == "hosted-gateway" { + defaultPath = hostedGatewayConfigPath + } + path := flags.String("config", defaultPath, "") + if flags.Parse(args) != nil || flags.NArg() != 0 || *protocol != mode { + return true, errHostedInvalid + } + ctx, stop := signal.NotifyContext(context.Background(), os.Interrupt, syscall.SIGTERM) + defer stop() + if mode == "hosted-gateway" { + settings, err := loadHostedGatewaySettings(*path, os.Getenv) + if err != nil { + return true, err + } + provider, err := newAzureFoundryTokenProvider() + if err != nil { + clear(settings.signingKey) + return true, errHostedInvalid + } + return true, serveHostedGateway(ctx, settings, provider) + } + cfg, err := loadHostedImageConfig(*path) + if err != nil { + return true, err + } + port := firstNonBlank(os.Getenv("PORT"), "8088") + n, err := strconv.Atoi(port) + // The supervisor and reverse relays bind these loopback ports after bootstrap. + if err != nil || n < 1 || n > 65535 || n == 8080 || n == 8091 || n == 8092 { + return true, errHostedInvalid + } + server, err := newHostedServer(ctx, cfg, os.Getenv, launchHostedSupervisor) + if err != nil { + return true, err + } + defer server.cancel() + return true, serveHostedHTTP(server.ctx, ":"+port, server) +} + +func serveHostedGateway(ctx context.Context, settings hostedGatewaySettings, provider foundryTokenProvider) error { + defer clear(settings.signingKey) + // Retain the listener before creating a one-shot remote lifetime. A local + // bind failure must not reserve a session or expose bootstrap credentials. + listener, err := net.Listen("tcp", settings.address) + if err != nil { + return errHostedInvalid + } + defer listener.Close() //nolint:errcheck + gateway, err := newHostedGateway(ctx, settings, provider) + if err != nil { + return err + } + defer gateway.close() + return serveHostedHTTPListener(gateway.ctx, listener, gateway) +} + +func serveHostedHTTP(ctx context.Context, address string, handler http.Handler) (serveErr error) { + if hosted, ok := handler.(*hostedServer); ok { + defer func() { serveErr = errors.Join(serveErr, hosted.close()) }() + } + listener, err := net.Listen("tcp", address) + if err != nil { + return errHostedInvalid + } + defer listener.Close() //nolint:errcheck + return serveHostedHTTPListener(ctx, listener, handler) +} + +func serveHostedHTTPListener(ctx context.Context, listener net.Listener, handler http.Handler) error { + server := newHostedHTTPServer(handler) + defer server.Close() //nolint:errcheck + done := make(chan error, 1) + go func() { done <- server.Serve(listener) }() + select { + case err := <-done: + if errors.Is(err, http.ErrServerClosed) { + return nil + } + return errHostedInvalid + case <-ctx.Done(): + // Close active requests too. Their acceptance is unknown to callers; + // this process never reconnects or creates another hosted lifetime. + _ = server.Close() + return nil + } +} + +func newHostedHTTPServer(handler http.Handler) *http.Server { + // Operation owners enforce byte limits and authorization/deadlines. A + // whole-body read timeout here truncates legitimate artifact uploads; + // header and idle timeouts bound inactive HTTP connections instead. + return &http.Server{Handler: handler, ReadHeaderTimeout: 10 * time.Second, + IdleTimeout: 2 * time.Minute, MaxHeaderBytes: 32 << 10} +} diff --git a/hosted_observation_test.go b/hosted_observation_test.go new file mode 100644 index 0000000..52bf459 --- /dev/null +++ b/hosted_observation_test.go @@ -0,0 +1,166 @@ +package main + +import ( + "bytes" + "errors" + "log" + "regexp" + "strconv" + "strings" + "sync" + "testing" + "time" + + "github.com/gorilla/websocket" +) + +const hostedTestCloseText = "synthetic-peer-text https://fixture.invalid/private?token=synthetic-secret" + +type hostedTestChannelLog struct { + mu sync.Mutex + data bytes.Buffer + logger *log.Logger +} + +func newHostedTestChannelLog() *hostedTestChannelLog { + l := &hostedTestChannelLog{} + l.logger = log.New(l, "", 0) + return l +} + +func (l *hostedTestChannelLog) Write(data []byte) (int, error) { + l.mu.Lock() + defer l.mu.Unlock() + return l.data.Write(data) +} + +func (l *hostedTestChannelLog) text() string { + l.mu.Lock() + defer l.mu.Unlock() + return l.data.String() +} + +type hostedTestChannelRecord struct { + role string + openedAt time.Time + duration time.Duration + code int +} + +func (l *hostedTestChannelLog) records(t *testing.T) []hostedTestChannelRecord { + t.Helper() + data := l.text() + if strings.Contains(data, hostedTestCloseText) || strings.Contains(data, "synthetic-secret") { + t.Fatal("channel lifecycle log exposed peer text") + } + if data == "" { + return nil + } + pattern := regexp.MustCompile(`^Foundry hosted channel closed role=(forward|reverse) opened_at=([^ ]+) duration_ms=([0-9]+) close_code=([0-9]+)$`) + var records []hostedTestChannelRecord + for _, line := range strings.Split(strings.TrimSuffix(data, "\n"), "\n") { + fields := pattern.FindStringSubmatch(line) + if fields == nil { + t.Fatal("channel lifecycle log contains fields outside the safe schema") + } + openedAt, timeErr := time.Parse(time.RFC3339Nano, fields[2]) + duration, durationErr := strconv.ParseInt(fields[3], 10, 64) + code, codeErr := strconv.Atoi(fields[4]) + if timeErr != nil || durationErr != nil || codeErr != nil { + t.Fatal("channel lifecycle log contains invalid safe metadata") + } + records = append(records, hostedTestChannelRecord{fields[1], openedAt, time.Duration(duration) * time.Millisecond, code}) + } + return records +} + +func hostedTestGoingAway(t *testing.T, ws *websocket.Conn) { + t.Helper() + if ws.WriteControl(websocket.CloseMessage, websocket.FormatCloseMessage(websocket.CloseGoingAway, hostedTestCloseText), time.Now().Add(time.Second)) != nil { + t.Fatal("could not send synthetic platform close code 1001") + } +} + +func hostedTestClosedChannelLogs(t *testing.T, logs *hostedTestChannelLog, affectedRole string, earliest, latest time.Time) { + t.Helper() + hostedServerWait(t, func() bool { return len(logs.records(t)) == 2 }) + seen := map[string]bool{} + for _, record := range logs.records(t) { + if seen[record.role] || record.openedAt.Before(earliest) || record.openedAt.After(latest) || + record.duration > time.Since(record.openedAt) { + t.Fatal("channel summary duplicated a role or lost its actual upgrade time") + } + seen[record.role] = true + if record.role == affectedRole && record.code != websocket.CloseGoingAway { + t.Fatal("affected channel did not retain observed close code 1001") + } + } + if !seen["forward"] || !seen["reverse"] { + t.Fatal("closed lifetime did not summarize both roles") + } +} + +func TestHostedWSConnCloseSummaryPreservesHandlerAndExcludesPeerText(t *testing.T) { + client, peer := hostedTestWSPair(t, false, nil) + logs := newHostedTestChannelLog() + openedAt := time.Now().Add(-time.Second) + previousCalled := false + previous := client.CloseHandler() + client.SetCloseHandler(func(code int, text string) error { + previousCalled = code == websocket.CloseGoingAway && text == hostedTestCloseText + return previous(code, text) + }) + observation := observeHostedChannel(client, "forward", openedAt, logs.logger) + conn := newHostedObservedWSConn(client, observation) + defer conn.Close() + hostedTestGoingAway(t, peer) + _, err := conn.Read(make([]byte, 1)) + var closed *websocket.CloseError + if !errors.As(err, &closed) || closed.Code != websocket.CloseGoingAway || closed.Text != hostedTestCloseText || !previousCalled { + t.Fatal("observation changed the original close handler or returned error") + } + _ = peer.SetReadDeadline(time.Now().Add(time.Second)) + _, _, err = peer.ReadMessage() + if !errors.As(err, &closed) || closed.Code != websocket.CloseGoingAway { + t.Fatal("observation changed Gorilla's close reply") + } + _ = conn.Close() + observation.finish() + records := logs.records(t) + if len(records) != 1 || records[0].role != "forward" || records[0].code != websocket.CloseGoingAway || + !records[0].openedAt.Equal(openedAt) || records[0].duration < time.Second { + t.Fatal("channel summary lost its original time/code or was emitted more than once") + } +} + +func TestHostedWSConnCloseSummaryDistinguishesUnobservedAndAbnormal(t *testing.T) { + for _, abrupt := range []bool{false, true} { + name := "local close without observed code" + if abrupt { + name = "peer EOF without close frame" + } + t.Run(name, func(t *testing.T) { + client, peer := hostedTestWSPair(t, false, nil) + logs := newHostedTestChannelLog() + observation := observeHostedChannel(client, "reverse", time.Now(), logs.logger) + conn := newHostedObservedWSConn(client, observation) + want := 0 + if abrupt { + _ = peer.Close() + _, _ = conn.Read(make([]byte, 1)) + want = websocket.CloseAbnormalClosure + } else { + observation.observeError(errors.New(hostedTestCloseText)) + } + var closed sync.WaitGroup + for range 8 { + closed.Go(func() { _ = conn.Close() }) + } + closed.Wait() + records := logs.records(t) + if len(records) != 1 || records[0].code != want { + t.Fatal("channel summary fabricated a code or duplicated concurrent closure") + } + }) + } +} diff --git a/hosted_process_linux.go b/hosted_process_linux.go new file mode 100644 index 0000000..6b66a5e --- /dev/null +++ b/hosted_process_linux.go @@ -0,0 +1,109 @@ +//go:build linux + +package main + +import ( + "context" + "os/exec" + "sort" + "syscall" + "time" + + "golang.org/x/sys/unix" +) + +func launchHostedSupervisor(ctx context.Context, environment map[string]string) (<-chan error, error) { + command := exec.Command("/usr/local/bin/orka-acp-runtime") + for name, value := range environment { + command.Env = append(command.Env, name+"="+value) + } + sort.Strings(command.Env) + return startHostedSupervisorProcess(ctx, command, hostedSupervisorStopGrace, hostedSupervisorKillWait) +} + +func startHostedSupervisorProcess(ctx context.Context, command *exec.Cmd, grace, killWait time.Duration) (<-chan error, error) { + if ctx.Err() != nil || grace <= 0 || killWait <= 0 { + return nil, errHostedInvalid + } + command.SysProcAttr = &syscall.SysProcAttr{Pdeathsig: syscall.SIGKILL, Setpgid: true} + // Nil output uses /dev/null directly, avoiding inherited copy pipes that + // could prevent Wait from returning after the supervisor has exited. + command.Stdout, command.Stderr = nil, nil + if command.Start() != nil { + return nil, errHostedInvalid + } + done := make(chan error, 1) + go func() { + done <- waitHostedSupervisorProcess(ctx, command, grace, killWait) + close(done) + }() + return done, nil +} + +func waitHostedSupervisorProcess(ctx context.Context, command *exec.Cmd, grace, killWait time.Duration) error { + exited := make(chan error, 1) + go func() { + var info unix.Siginfo + for { + err := unix.Waitid(unix.P_PID, command.Process.Pid, &info, unix.WEXITED|unix.WNOWAIT, nil) + if err != syscall.EINTR { + exited <- err + return + } + } + }() + // Keep the leader unreaped until every group signal is sent. Otherwise a + // fast exit and PID reuse could redirect a later kill(-pgid) elsewhere. + var observedErr, signalErr error + forced := false + select { + case observedErr = <-exited: + case <-ctx.Done(): + signalErr = signalHostedSupervisorGroup(command.Process.Pid, syscall.SIGTERM) + graceTimer := time.NewTimer(grace) + select { + case observedErr = <-exited: + case <-graceTimer.C: + forced = true + if err := signalHostedSupervisorGroup(command.Process.Pid, syscall.SIGKILL); err != nil { + signalErr = err + } + killTimer := time.NewTimer(killWait) + select { + case observedErr = <-exited: + case <-killTimer.C: + // A stuck kernel task is not cleanup proof. Reap if it later + // exits, but let the caller retire this failed lifetime now. + go func() { <-exited; _ = command.Wait() }() + return errHostedInvalid + } + killTimer.Stop() + } + graceTimer.Stop() + } + if observedErr != nil { + _ = command.Process.Kill() + go func() { _ = command.Wait() }() + return errHostedInvalid + } + // This reaches only the supervisor's own group. ACP children deliberately + // use separate groups and UIDs: their cleanup belongs to Orka's bounded + // Server.Close, and remote cleanup still requires its broker receipts. + if err := signalHostedSupervisorGroup(command.Process.Pid, syscall.SIGKILL); err != nil { + signalErr = err + } + if err := command.Wait(); err != nil || signalErr != nil || forced { + return errHostedInvalid + } + return nil +} + +func signalHostedSupervisorGroup(pid int, signal syscall.Signal) error { + if pid <= 0 { + return errHostedInvalid + } + if err := syscall.Kill(-pid, signal); err != nil && err != syscall.ESRCH { + return err + } + return nil +} diff --git a/hosted_process_linux_test.go b/hosted_process_linux_test.go new file mode 100644 index 0000000..cda3ad5 --- /dev/null +++ b/hosted_process_linux_test.go @@ -0,0 +1,180 @@ +//go:build linux + +package main + +import ( + "context" + "errors" + "os" + "os/exec" + "os/signal" + "path/filepath" + "strconv" + "syscall" + "testing" + "time" + + "golang.org/x/sys/unix" +) + +func TestHostedSupervisorProcessShutdownScopes(t *testing.T) { + for _, mode := range []string{"graceful", "force-group", "exit-group", "force-detached"} { + t.Run(mode, func(t *testing.T) { + directory := t.TempDir() + command := hostedProcessTestCommand(mode, directory) + ctx, cancel := context.WithCancel(t.Context()) + defer cancel() + done, err := startHostedSupervisorProcess(ctx, command, 250*time.Millisecond, 2*time.Second) + if err != nil { + t.Fatal("could not launch the isolated supervisor fixture") + } + t.Cleanup(func() { + cancel() + _ = command.Process.Kill() + select { + case <-done: + case <-time.After(3 * time.Second): + t.Error("supervisor fixture did not finish cleanup") + } + }) + hostedProcessTestFile(t, filepath.Join(directory, "ready")) + childFD := -1 + if mode != "graceful" { + childPID, err := strconv.Atoi(string(hostedProcessTestFile(t, filepath.Join(directory, "child-pid")))) + if err != nil || childPID <= 0 { + t.Fatal("fixture did not identify its child") + } + childFD, err = unix.PidfdOpen(childPID, 0) + if err != nil { + t.Fatal("could not retain the exact fixture child") + } + t.Cleanup(func() { + _ = unix.PidfdSendSignal(childFD, unix.SIGKILL, nil, 0) + hostedProcessTestExited(t, childFD, true) + _ = unix.Close(childFD) + }) + } + cancel() + select { + case err := <-done: + forced := mode == "force-group" || mode == "force-detached" + if forced && !errors.Is(err, errHostedInvalid) || !forced && err != nil { + t.Fatal("supervisor exit was classified incorrectly") + } + status, ok := command.ProcessState.Sys().(syscall.WaitStatus) + if !ok || forced && (!status.Signaled() || status.Signal() != syscall.SIGKILL) || !forced && status.ExitStatus() != 0 { + t.Fatal("supervisor kernel exit did not match the requested shutdown") + } + case <-time.After(3 * time.Second): + t.Fatal("supervisor shutdown exceeded the fixture bound") + } + if mode == "graceful" { + hostedProcessTestFile(t, filepath.Join(directory, "cleaned")) + } else { + // An independent ACP process group is intentionally outside + // this fallback. Only Orka's UID owner can prove its cleanup. + hostedProcessTestExited(t, childFD, mode != "force-detached") + } + }) + } +} + +func TestHostedSupervisorProcessRejectsCancelledLaunch(t *testing.T) { + ctx, cancel := context.WithCancel(t.Context()) + cancel() + command := hostedProcessTestCommand("graceful", t.TempDir()) + if _, err := startHostedSupervisorProcess(ctx, command, time.Second, time.Second); err == nil || command.Process != nil { + t.Fatal("cancelled hosted lifetime launched a process") + } + for _, pid := range []int{0, -1} { + if signalHostedSupervisorGroup(pid, syscall.SIGKILL) == nil { + t.Fatal("invalid supervisor identity was accepted for group signaling") + } + } +} + +func hostedProcessTestCommand(mode, directory string) *exec.Cmd { + command := exec.Command(os.Args[0], "-test.run=^TestHostedProcessHelper$") + // A race-instrumented fixture must not add the detector's default one-second + // exit sleep to the deliberately short graceful-shutdown test window. + command.Env = []string{"FOUNDRY_HOSTED_PROCESS_HELPER=" + mode, "FOUNDRY_HOSTED_PROCESS_DIRECTORY=" + directory, + "GORACE=atexit_sleep_ms=0"} + return command +} + +func hostedProcessTestFile(t *testing.T, path string) []byte { + t.Helper() + deadline := time.NewTimer(3 * time.Second) + defer deadline.Stop() + tick := time.NewTicker(5 * time.Millisecond) + defer tick.Stop() + for { + if data, err := os.ReadFile(path); err == nil && len(data) != 0 { + return data + } + select { + case <-deadline.C: + t.Fatal("supervisor fixture did not reach its checkpoint") + case <-tick.C: + } + } +} + +func hostedProcessTestExited(t *testing.T, fd int, want bool) { + t.Helper() + timeout := 0 + if want { + timeout = 1000 + } + fds := []unix.PollFd{{Fd: int32(fd), Events: unix.POLLIN}} + n, err := unix.Poll(fds, timeout) + if err != nil || (n > 0 && fds[0].Revents&unix.POLLIN != 0) != want { + t.Fatal("exact fixture child exit state did not match its process-group scope") + } +} + +func TestHostedProcessHelper(t *testing.T) { + mode := os.Getenv("FOUNDRY_HOSTED_PROCESS_HELPER") + if mode == "" { + return + } + directory := os.Getenv("FOUNDRY_HOSTED_PROCESS_DIRECTORY") + write := func(name, value string) { + if os.WriteFile(filepath.Join(directory, name), []byte(value), 0o600) != nil { + os.Exit(3) + } + } + if mode == "child" { + signal.Ignore(syscall.SIGTERM) + write("child-pid", strconv.Itoa(os.Getpid())) + for { + time.Sleep(time.Hour) + } + } + terminated := make(chan os.Signal, 1) + if mode == "force-group" || mode == "force-detached" { + signal.Ignore(syscall.SIGTERM) + } else { + signal.Notify(terminated, syscall.SIGTERM) + } + if mode != "graceful" { + child := hostedProcessTestCommand("child", directory) + if mode == "force-detached" { + child.SysProcAttr = &syscall.SysProcAttr{Setpgid: true} + } + if child.Start() != nil { + os.Exit(4) + } + } + write("ready", "1") + if mode == "force-group" || mode == "force-detached" { + for { + time.Sleep(time.Hour) + } + } + <-terminated + if mode == "graceful" { + time.Sleep(25 * time.Millisecond) + write("cleaned", "1") + } +} diff --git a/hosted_process_other.go b/hosted_process_other.go new file mode 100644 index 0000000..58cf063 --- /dev/null +++ b/hosted_process_other.go @@ -0,0 +1,9 @@ +//go:build !linux + +package main + +import "context" + +func launchHostedSupervisor(context.Context, map[string]string) (<-chan error, error) { + return nil, errHostedInvalid +} diff --git a/hosted_protocol.go b/hosted_protocol.go new file mode 100644 index 0000000..648897e --- /dev/null +++ b/hosted_protocol.go @@ -0,0 +1,298 @@ +package main + +import ( + "crypto/ed25519" + "crypto/subtle" + "encoding/base64" + "encoding/json" + "errors" + "net/url" + "strconv" + "strings" + "time" + "unicode" + "unicode/utf8" + + "filippo.io/edwards25519" + "github.com/google/uuid" +) + +const hostedProtocol = "orka.foundry.hosted.v1" + +var errHostedInvalid = errors.New("invalid Foundry hosted configuration or handshake") + +type hostedImageConfig struct { + Protocol string `json:"protocol"` + DeploymentID string `json:"deploymentID"` + Target acpHostedTarget `json:"target"` + SigningPublicKey string `json:"signingPublicKey"` + AgentConfigurationDigest string `json:"agentConfigurationDigest"` +} + +type hostedBootstrap struct { + Environment map[string]string `json:"environment"` + ControllerToken string `json:"controllerToken"` + CapabilitySecret string `json:"capabilitySecret"` + ProviderToken string `json:"providerToken"` +} + +type hostedChallenge struct { + Protocol string `json:"protocol"` + DeploymentID string `json:"deploymentID"` + ConfigurationDigest string `json:"configurationDigest"` + AgentName string `json:"agentName"` + AgentVersion string `json:"agentVersion"` + SessionID string `json:"sessionID"` + BootID string `json:"bootID"` + Nonce string `json:"nonce"` +} + +type hostedHello struct { + Challenge hostedChallenge `json:"challenge"` + PairID string `json:"pairID"` + Role string `json:"role"` + BootstrapDigest string `json:"bootstrapDigest"` + ExpiresAt int64 `json:"expiresAt"` + Signature string `json:"signature"` +} + +type hostedAccepted struct { + Protocol string `json:"protocol"` + PairID string `json:"pairID"` + Role string `json:"role"` + BootID string `json:"bootID"` + BootstrapDigest string `json:"bootstrapDigest"` +} + +func validateHostedImageConfig(cfg hostedImageConfig) error { + if cfg.Protocol != hostedProtocol || !hostedUUIDValid(cfg.DeploymentID) || + !brokerDigestValid(cfg.AgentConfigurationDigest) || !hostedTargetValid(cfg.Target) { + return errHostedInvalid + } + if _, ok := hostedPublicKey(cfg.SigningPublicKey); !ok { + return errHostedInvalid + } + return nil +} + +func hostedSigningKeyValid(key ed25519.PrivateKey) bool { + if len(key) != ed25519.PrivateKeySize || !hostedNonzeroBytes(key[:ed25519.SeedSize]) { + return false + } + derived := ed25519.NewKeyFromSeed(key[:ed25519.SeedSize]) + defer clear(derived) + return subtle.ConstantTimeCompare(key, derived) == 1 +} + +func signHostedHello(value hostedHello, key ed25519.PrivateKey) (hostedHello, error) { + if !hostedSigningKeyValid(key) { + return hostedHello{}, errHostedInvalid + } + data, err := hostedHelloSigningBytes(value) + if err != nil { + return hostedHello{}, err + } + value.Signature = base64.RawURLEncoding.EncodeToString(ed25519.Sign(key, data)) + return value, nil +} + +// Verification authenticates one exact challenge and its bootstrap digest. +// The caller must separately consume the nonce and enforce pair/role ownership. +func verifyHostedHello(value hostedHello, expected hostedChallenge, publicKey string, now time.Time) error { + if value.Challenge != expected { + return errHostedInvalid + } + expires := time.Unix(value.ExpiresAt, 0) + if !expires.After(now) || expires.After(now.Add(90*time.Second)) { + return errHostedInvalid + } + key, keyOK := hostedPublicKey(publicKey) + signature, signatureOK := hostedCanonicalBytes(value.Signature, ed25519.SignatureSize) + data, err := hostedHelloSigningBytes(value) + if !keyOK || !signatureOK || err != nil || !ed25519.Verify(key, data, signature) { + return errHostedInvalid + } + return nil +} + +func hostedHelloSigningBytes(value hostedHello) ([]byte, error) { + c := value.Challenge + if c.Protocol != hostedProtocol || !hostedUUIDValid(c.DeploymentID) || + !brokerDigestValid(c.ConfigurationDigest) || validateAgentName(c.AgentName) != nil || + !hostedPositiveUint(c.AgentVersion) || !hostedUUIDValid(c.SessionID) || !hostedUUIDValid(c.BootID) || + !hostedUUIDValid(value.PairID) || (value.Role != "forward" && value.Role != "reverse") || + !brokerDigestValid(value.BootstrapDigest) || value.ExpiresAt <= 0 { + return nil, errHostedInvalid + } + if _, ok := hostedCanonicalBytes(c.Nonce, 32); !ok { + return nil, errHostedInvalid + } + // Keep the field order stable. Neither input JSON order nor the signature + // itself participates in this domain-separated signing representation. + unsigned := struct { + Challenge hostedChallenge `json:"challenge"` + PairID string `json:"pairID"` + Role string `json:"role"` + BootstrapDigest string `json:"bootstrapDigest"` + ExpiresAt int64 `json:"expiresAt"` + }{value.Challenge, value.PairID, value.Role, value.BootstrapDigest, value.ExpiresAt} + data, err := json.Marshal(unsigned) + if err != nil { + return nil, errHostedInvalid + } + return append([]byte(hostedProtocol+"\x00hello\x00"), data...), nil +} + +func validateHostedBootstrap(cfg hostedImageConfig, value hostedBootstrap) error { + if validateHostedImageConfig(cfg) != nil || !hostedCredentialValid(value.ControllerToken) || + !hostedCredentialValid(value.CapabilitySecret) || !hostedCredentialValid(value.ProviderToken) { + return errHostedInvalid + } + required := 0 + for key, value := range value.Environment { + valid := false + switch key { + case "ORKA_ACP_PROVIDER": + valid = value == "foundry" + case "ORKA_ACP_MODEL": + valid = utf8.ValidString(value) && acpSafeString(value, 512) && strings.TrimSpace(value) == value + case "ORKA_ACP_FOUNDRY_ADAPTER_DIGEST", "ORKA_ACP_TOOL_POLICY_DIGEST", + "ORKA_ACP_APPROVAL_POLICY_DIGEST", "ORKA_ACP_MCP_CONFIGURATION_DIGEST": + valid = brokerDigestValid(value) + case "ORKA_ACP_AGENT_CONFIGURATION_DIGEST": + valid = value == cfg.AgentConfigurationDigest + case "ORKA_ACP_WORKSPACE_INTENT": + valid = value == "read" || value == "write" + case "ORKA_ACP_PROXY_CREDENTIAL_ROLE": + valid = value == "operator-managed" + case "ORKA_ACP_PROXY_CREDENTIAL_SCOPE": + valid = value == "external-runtime" + case "ORKA_ACP_RESOURCE_CLASS": + valid = value == "external" + case "ORKA_ACP_TRUST_NAMESPACE": + valid = hostedDNSLabelValid(value) + case "ORKA_ACP_CONTROLLER_EPOCH", "ORKA_ACP_RUNTIME_POOL_GENERATION": + valid = hostedPositiveUint(value) + case "ORKA_ACP_RUNTIME_POOL_UID": + valid = hostedUUIDValid(value) + case "ORKA_ACP_MODEL_CONTEXT_LIMIT", "ORKA_ACP_MODEL_OUTPUT_LIMIT": + if !hostedPositiveUint(value) { + return errHostedInvalid + } + continue + default: + return errHostedInvalid + } + if !valid { + return errHostedInvalid + } + required++ + } + _, contextLimit := value.Environment["ORKA_ACP_MODEL_CONTEXT_LIMIT"] + _, outputLimit := value.Environment["ORKA_ACP_MODEL_OUTPUT_LIMIT"] + if required != 15 || contextLimit != outputLimit { + return errHostedInvalid + } + return nil +} + +func hostedTargetValid(target acpHostedTarget) bool { + if validateAgentName(target.AgentName) != nil || !hostedPositiveUint(target.AgentVersion) || + !acpSafeString(target.ProjectEndpoint, 2048) || strings.ContainsAny(target.ProjectEndpoint, "%#") { + return false + } + u, err := url.Parse(target.ProjectEndpoint) + if err != nil || u.Scheme != "https" || u.Opaque != "" || u.User != nil || u.Host != u.Hostname() || + u.RawQuery != "" || u.ForceQuery || u.Fragment != "" || u.RawPath != "" || u.String() != target.ProjectEndpoint { + return false + } + const suffix = ".services.ai.azure.com" + if !strings.HasSuffix(u.Host, suffix) || !hostedDNSLabelValid(strings.TrimSuffix(u.Host, suffix)) { + return false + } + const prefix = "/api/projects/" + if !strings.HasPrefix(u.Path, prefix) { + return false + } + project := strings.TrimPrefix(u.Path, prefix) + if len(project) == 0 || len(project) > 128 || !hostedASCIIAlphanumeric(project[0]) || + !hostedASCIIAlphanumeric(project[len(project)-1]) { + return false + } + for i := range len(project) { + if !hostedASCIIAlphanumeric(project[i]) && project[i] != '-' && project[i] != '_' && project[i] != '.' { + return false + } + } + return true +} + +func hostedUUIDValid(value string) bool { + id, err := uuid.Parse(value) + return err == nil && id != uuid.Nil && id.String() == value +} + +func hostedPositiveUint(value string) bool { + number, err := strconv.ParseUint(value, 10, 64) + return err == nil && number > 0 && strconv.FormatUint(number, 10) == value +} + +func hostedCanonicalBytes(value string, size int) ([]byte, bool) { + if len(value) != base64.RawURLEncoding.EncodedLen(size) { + return nil, false + } + decoded, err := base64.RawURLEncoding.Strict().DecodeString(value) + if err != nil || len(decoded) != size || base64.RawURLEncoding.EncodeToString(decoded) != value || + !hostedNonzeroBytes(decoded) { + return nil, false + } + return decoded, true +} + +func hostedPublicKey(value string) (ed25519.PublicKey, bool) { + decoded, ok := hostedCanonicalBytes(value, ed25519.PublicKeySize) + if !ok { + return nil, false + } + point, err := new(edwards25519.Point).SetBytes(decoded) + if err != nil || subtle.ConstantTimeCompare(point.Bytes(), decoded) != 1 { + return nil, false + } + // A small-order public key can verify a signature without a private key. + if new(edwards25519.Point).MultByCofactor(point).Equal(edwards25519.NewIdentityPoint()) == 1 { + return nil, false + } + return ed25519.PublicKey(decoded), true +} + +func hostedNonzeroBytes(value []byte) bool { + var combined byte + for _, b := range value { + combined |= b + } + return combined != 0 +} + +func hostedCredentialValid(value string) bool { + if len(value) < 32 || !utf8.ValidString(value) || !acpSafeString(value, 16<<10) { + return false + } + return strings.IndexFunc(value, unicode.IsSpace) < 0 +} + +func hostedDNSLabelValid(value string) bool { + if len(value) == 0 || len(value) > 63 || value != strings.ToLower(value) || + !hostedASCIIAlphanumeric(value[0]) || !hostedASCIIAlphanumeric(value[len(value)-1]) { + return false + } + for i := range len(value) { + if !hostedASCIIAlphanumeric(value[i]) && value[i] != '-' { + return false + } + } + return true +} + +func hostedASCIIAlphanumeric(value byte) bool { + return value >= 'a' && value <= 'z' || value >= 'A' && value <= 'Z' || value >= '0' && value <= '9' +} diff --git a/hosted_protocol_test.go b/hosted_protocol_test.go new file mode 100644 index 0000000..c3c3305 --- /dev/null +++ b/hosted_protocol_test.go @@ -0,0 +1,559 @@ +package main + +import ( + "bytes" + "crypto/ed25519" + "crypto/rand" + "encoding/base64" + "encoding/json" + "maps" + "strings" + "testing" + "time" + + "github.com/google/uuid" +) + +type hostedProtocolFixture struct { + config hostedImageConfig + bootstrap hostedBootstrap + hello hostedHello + key ed25519.PrivateKey + now time.Time +} + +func newHostedProtocolFixture(t *testing.T) hostedProtocolFixture { + t.Helper() + public, private, err := ed25519.GenerateKey(rand.Reader) + if err != nil { + t.Fatal("could not generate test signing key") + } + cfg := hostedImageConfig{ + Protocol: hostedProtocol, DeploymentID: uuid.NewString(), + Target: acpHostedTarget{ + ProjectEndpoint: "https://test-account.services.ai.azure.com/api/projects/test-project", + AgentName: "test-agent", + AgentVersion: "8", + }, + SigningPublicKey: base64.RawURLEncoding.EncodeToString(public), + AgentConfigurationDigest: brokerSHA([]byte("test agent configuration")), + } + bootstrap := hostedBootstrap{ + ControllerToken: strings.Repeat("test-controller-", 3), + CapabilitySecret: strings.Repeat("test-capability-", 3), + ProviderToken: strings.Repeat("test-provider-", 3), + Environment: map[string]string{ + "ORKA_ACP_PROVIDER": "foundry", + "ORKA_ACP_MODEL": "test-model", + "ORKA_ACP_FOUNDRY_ADAPTER_DIGEST": brokerSHA([]byte("test adapter")), + "ORKA_ACP_WORKSPACE_INTENT": "read", + "ORKA_ACP_AGENT_CONFIGURATION_DIGEST": cfg.AgentConfigurationDigest, + "ORKA_ACP_TOOL_POLICY_DIGEST": brokerSHA([]byte("test tool policy")), + "ORKA_ACP_APPROVAL_POLICY_DIGEST": brokerSHA([]byte("test approval policy")), + "ORKA_ACP_MCP_CONFIGURATION_DIGEST": brokerSHA([]byte("test MCP configuration")), + "ORKA_ACP_PROXY_CREDENTIAL_ROLE": "operator-managed", + "ORKA_ACP_PROXY_CREDENTIAL_SCOPE": "external-runtime", + "ORKA_ACP_RESOURCE_CLASS": "external", + "ORKA_ACP_TRUST_NAMESPACE": "test-namespace", + "ORKA_ACP_CONTROLLER_EPOCH": "1", + "ORKA_ACP_RUNTIME_POOL_GENERATION": "2", + "ORKA_ACP_RUNTIME_POOL_UID": uuid.NewString(), + }, + } + now := time.Unix(1_800_000_000, 0) + hello := hostedHello{ + Challenge: hostedChallenge{ + Protocol: hostedProtocol, DeploymentID: cfg.DeploymentID, + ConfigurationDigest: brokerJSONDigest(cfg), AgentName: cfg.Target.AgentName, + AgentVersion: cfg.Target.AgentVersion, SessionID: uuid.NewString(), BootID: uuid.NewString(), + Nonce: base64.RawURLEncoding.EncodeToString(bytes.Repeat([]byte{7}, 32)), + }, + PairID: uuid.NewString(), Role: "forward", BootstrapDigest: brokerJSONDigest(bootstrap), + ExpiresAt: now.Add(time.Minute).Unix(), + } + return hostedProtocolFixture{config: cfg, bootstrap: bootstrap, hello: hello, key: private, now: now} +} + +func (f hostedProtocolFixture) sign(t *testing.T, value hostedHello) hostedHello { + t.Helper() + signed, err := signHostedHello(value, f.key) + if err != nil { + t.Fatal("valid test handshake was not signed") + } + return signed +} + +func TestHostedHelloAuthenticatesBothRoles(t *testing.T) { + f := newHostedProtocolFixture(t) + for _, role := range []string{"forward", "reverse"} { + t.Run(role, func(t *testing.T) { + value := f.hello + value.Role = role + signed := f.sign(t, value) + if verifyHostedHello(signed, value.Challenge, f.config.SigningPublicKey, f.now) != nil { + t.Fatal("valid signed role rejected") + } + // Verification is stateless. The hosted server consumes each role and + // nonce under its own admission lock after cryptographic verification. + if verifyHostedHello(signed, value.Challenge, f.config.SigningPublicKey, f.now) != nil { + t.Fatal("signature verification unexpectedly consumed handshake state") + } + }) + } +} + +func TestHostedHelloBindsEverySignedField(t *testing.T) { + f := newHostedProtocolFixture(t) + signed := f.sign(t, f.hello) + for name, mutate := range map[string]func(*hostedHello){ + "protocol": func(v *hostedHello) { v.Challenge.Protocol = "other-protocol" }, + "deployment": func(v *hostedHello) { v.Challenge.DeploymentID = uuid.NewString() }, + "configuration": func(v *hostedHello) { v.Challenge.ConfigurationDigest = brokerSHA([]byte("other configuration")) }, + "agent": func(v *hostedHello) { v.Challenge.AgentName = "other-agent" }, + "version": func(v *hostedHello) { v.Challenge.AgentVersion = "9" }, + "session": func(v *hostedHello) { v.Challenge.SessionID = uuid.NewString() }, + "boot": func(v *hostedHello) { v.Challenge.BootID = uuid.NewString() }, + "nonce": func(v *hostedHello) { + v.Challenge.Nonce = base64.RawURLEncoding.EncodeToString(bytes.Repeat([]byte{8}, 32)) + }, + "pair": func(v *hostedHello) { v.PairID = uuid.NewString() }, + "role": func(v *hostedHello) { v.Role = "reverse" }, + "expiry": func(v *hostedHello) { v.ExpiresAt-- }, + "bootstrap body": func(v *hostedHello) { + body := f.bootstrap + body.Environment = maps.Clone(body.Environment) + body.Environment["ORKA_ACP_WORKSPACE_INTENT"] = "write" + v.BootstrapDigest = brokerJSONDigest(body) + }, + "bootstrap credential": func(v *hostedHello) { + body := f.bootstrap + body.ControllerToken = strings.Repeat("different-test-controller-", 2) + v.BootstrapDigest = brokerJSONDigest(body) + }, + } { + t.Run(name, func(t *testing.T) { + changed := signed + mutate(&changed) + // Matching the modified challenge deliberately bypasses the equality + // check, so valid field mutations must fail the signature itself. + if verifyHostedHello(changed, changed.Challenge, f.config.SigningPublicKey, f.now) == nil { + t.Fatal("changed signed field accepted") + } + }) + } +} + +func TestHostedHelloRejectsDifferentExpectedChallenge(t *testing.T) { + f := newHostedProtocolFixture(t) + signed := f.sign(t, f.hello) + for name, mutate := range map[string]func(*hostedChallenge){ + "deployment": func(c *hostedChallenge) { c.DeploymentID = uuid.NewString() }, + "session": func(c *hostedChallenge) { c.SessionID = uuid.NewString() }, + "boot": func(c *hostedChallenge) { c.BootID = uuid.NewString() }, + "nonce": func(c *hostedChallenge) { + c.Nonce = base64.RawURLEncoding.EncodeToString(bytes.Repeat([]byte{9}, 32)) + }, + } { + t.Run(name, func(t *testing.T) { + expected := signed.Challenge + mutate(&expected) + if verifyHostedHello(signed, expected, f.config.SigningPublicKey, f.now) == nil { + t.Fatal("handshake authenticated against another challenge") + } + }) + } +} + +func TestHostedHelloExpiryWindow(t *testing.T) { + f := newHostedProtocolFixture(t) + for _, test := range []struct { + name string + offset time.Duration + valid bool + }{ + {"expired", -time.Second, false}, + {"exactly now", 0, false}, + {"one second", time.Second, true}, + {"ninety seconds", 90 * time.Second, true}, + {"too far ahead", 91 * time.Second, false}, + } { + t.Run(test.name, func(t *testing.T) { + value := f.hello + value.ExpiresAt = f.now.Add(test.offset).Unix() + signed := f.sign(t, value) + if (verifyHostedHello(signed, signed.Challenge, f.config.SigningPublicKey, f.now) == nil) != test.valid { + t.Fatal("expiry window enforced incorrectly") + } + }) + } + signed := f.sign(t, f.hello) + if verifyHostedHello(signed, signed.Challenge, f.config.SigningPublicKey, + time.Unix(signed.ExpiresAt, 1)) == nil { + t.Fatal("subsecond expiry was rounded into the acceptance window") + } +} + +func TestHostedHelloRejectsMalformedKeysAndSignatures(t *testing.T) { + f := newHostedProtocolFixture(t) + signed := f.sign(t, f.hello) + corruptedKey := bytes.Clone(f.key) + corruptedKey[len(corruptedKey)-1] ^= 1 + for name, key := range map[string]ed25519.PrivateKey{ + "missing": nil, "seed only": f.key[:ed25519.SeedSize], + "short": f.key[:len(f.key)-1], "long": append(bytes.Clone(f.key), 1), + "zero": make([]byte, ed25519.PrivateKeySize), + "zero seed": ed25519.NewKeyFromSeed(make([]byte, ed25519.SeedSize)), + "inconsistent public half": corruptedKey, + } { + t.Run("private "+name, func(t *testing.T) { + if _, err := signHostedHello(f.hello, key); err == nil { + t.Fatal("malformed private key accepted") + } + }) + } + other := newHostedProtocolFixture(t) + for name, key := range map[string]string{ + "missing": "", "different": other.config.SigningPublicKey, + "zero": base64.RawURLEncoding.EncodeToString(make([]byte, ed25519.PublicKeySize)), + "short": base64.RawURLEncoding.EncodeToString(bytes.Repeat([]byte{1}, ed25519.PublicKeySize-1)), + "long": base64.RawURLEncoding.EncodeToString(bytes.Repeat([]byte{1}, ed25519.PublicKeySize+1)), + "padded": f.config.SigningPublicKey + "=", + "newline": f.config.SigningPublicKey + "\n", + "tail bits": hostedTestNoncanonicalBase64(f.config.SigningPublicKey), + } { + t.Run("public "+name, func(t *testing.T) { + if verifyHostedHello(signed, signed.Challenge, key, f.now) == nil { + t.Fatal("invalid or different public key accepted") + } + }) + } + for name, signature := range map[string]string{ + "missing": "", "wrong alphabet": strings.Repeat("+", len(signed.Signature)), + "zero": base64.RawURLEncoding.EncodeToString(make([]byte, ed25519.SignatureSize)), + "short": base64.RawURLEncoding.EncodeToString(bytes.Repeat([]byte{1}, ed25519.SignatureSize-1)), + "long": base64.RawURLEncoding.EncodeToString(bytes.Repeat([]byte{1}, ed25519.SignatureSize+1)), + "padded": signed.Signature + "==", + "newline": signed.Signature + "\n", + "tail bits": hostedTestNoncanonicalBase64(signed.Signature), + } { + t.Run("signature "+name, func(t *testing.T) { + value := signed + value.Signature = signature + if verifyHostedHello(value, value.Challenge, f.config.SigningPublicKey, f.now) == nil { + t.Fatal("malformed signature accepted") + } + }) + } +} + +func hostedTestNoncanonicalBase64(value string) string { + const alphabet = "ABCDEFGHIJKLMNOPQRSTUVWXYZabcdefghijklmnopqrstuvwxyz0123456789-_" + last := strings.IndexByte(alphabet, value[len(value)-1]) + return value[:len(value)-1] + string(alphabet[last+1]) +} + +func TestHostedHelloRejectsIdentityPublicKey(t *testing.T) { + f := newHostedProtocolFixture(t) + identity := make([]byte, ed25519.PublicKeySize) + identity[0] = 1 + // For an identity public key, R = identity and S = 0 can satisfy the + // standard Ed25519 verification equation for every handshake body. + forged := make([]byte, ed25519.SignatureSize) + forged[0] = 1 + f.config.SigningPublicKey = base64.RawURLEncoding.EncodeToString(identity) + f.hello.Signature = base64.RawURLEncoding.EncodeToString(forged) + if validateHostedImageConfig(f.config) == nil { + t.Error("identity public key accepted in image configuration") + } + if verifyHostedHello(f.hello, f.hello.Challenge, f.config.SigningPublicKey, f.now) == nil { + t.Error("handshake authenticated without knowledge of a private key") + } +} + +func TestHostedImageConfigRejectsInvalidPublicPoints(t *testing.T) { + f := newHostedProtocolFixture(t) + offCurve := make([]byte, ed25519.PublicKeySize) + offCurve[0] = 2 + orderTwo := bytes.Repeat([]byte{0xff}, ed25519.PublicKeySize) + orderTwo[0], orderTwo[31] = 0xec, 0x7f // y = p - 1 + orderFour := make([]byte, ed25519.PublicKeySize) + orderFour[31] = 0x80 // y = 0, negative x + negativeIdentity := make([]byte, ed25519.PublicKeySize) + negativeIdentity[0], negativeIdentity[31] = 1, 0x80 + noncanonicalPoint := bytes.Repeat([]byte{0xff}, ed25519.PublicKeySize) + noncanonicalPoint[0], noncanonicalPoint[31] = 0xf0, 0x7f // y = p + 3 + for name, key := range map[string][]byte{ + "off curve": offCurve, "order two": orderTwo, "order four": orderFour, + "negative identity": negativeIdentity, "noncanonical point": noncanonicalPoint, + } { + t.Run(name, func(t *testing.T) { + cfg := f.config + cfg.SigningPublicKey = base64.RawURLEncoding.EncodeToString(key) + if validateHostedImageConfig(cfg) == nil { + t.Fatal("invalid curve point accepted as signing public key") + } + }) + } +} + +func TestHostedHelloRejectsInvalidSigningInputs(t *testing.T) { + f := newHostedProtocolFixture(t) + for name, mutate := range map[string]func(*hostedHello){ + "unknown protocol": func(v *hostedHello) { v.Challenge.Protocol = "other" }, + "invalid deployment": func(v *hostedHello) { + v.Challenge.DeploymentID = strings.ReplaceAll(v.Challenge.DeploymentID, "-", "") + }, + "invalid configuration": func(v *hostedHello) { v.Challenge.ConfigurationDigest = "invalid" }, + "invalid agent": func(v *hostedHello) { v.Challenge.AgentName = "../other" }, + "unpinned version": func(v *hostedHello) { v.Challenge.AgentVersion = "@latest" }, + "invalid session": func(v *hostedHello) { v.Challenge.SessionID = "session" }, + "invalid boot": func(v *hostedHello) { v.Challenge.BootID = "boot" }, + "zero nonce": func(v *hostedHello) { + v.Challenge.Nonce = base64.RawURLEncoding.EncodeToString(make([]byte, 32)) + }, + "padded nonce": func(v *hostedHello) { v.Challenge.Nonce += "=" }, + "invalid pair": func(v *hostedHello) { v.PairID = "pair" }, + "unknown role": func(v *hostedHello) { v.Role = "control" }, + "invalid bootstrap": func(v *hostedHello) { v.BootstrapDigest = "invalid" }, + "zero expiry": func(v *hostedHello) { v.ExpiresAt = 0 }, + "negative expiry": func(v *hostedHello) { v.ExpiresAt = -1 }, + } { + t.Run(name, func(t *testing.T) { + value := f.hello + mutate(&value) + if _, err := signHostedHello(value, f.key); err == nil { + t.Fatal("invalid handshake signed") + } + }) + } +} + +func TestHostedHelloCanonicalJSONWire(t *testing.T) { + f := newHostedProtocolFixture(t) + signed := f.sign(t, f.hello) + raw, err := json.Marshal(signed) + if err != nil { + t.Fatal("could not encode test handshake") + } + var reordered map[string]json.RawMessage + if acpDecode(raw, &reordered, true) != nil { + t.Fatal("could not decode test wire fields") + } + // Marshalling a map sorts the keys differently from the signing struct. + raw, err = json.Marshal(reordered) + if err != nil { + t.Fatal("could not reorder test wire fields") + } + var decoded hostedHello + if acpDecode(raw, &decoded, true) != nil || + verifyHostedHello(decoded, signed.Challenge, f.config.SigningPublicKey, f.now) != nil { + t.Fatal("valid reordered wire representation rejected") + } + reordered["unexpected"] = json.RawMessage(`true`) + raw, _ = json.Marshal(reordered) + if acpDecode(raw, &decoded, true) == nil { + t.Fatal("unexpected handshake field accepted by strict wire decoder") + } +} + +func TestHostedImageConfigRequiresExactTarget(t *testing.T) { + f := newHostedProtocolFixture(t) + if validateHostedImageConfig(f.config) != nil { + t.Fatal("valid hosted image configuration rejected") + } + for _, endpoint := range []string{ + "http://test-account.services.ai.azure.com/api/projects/test-project", + "https://test-account.services.ai.azure.com:443/api/projects/test-project", + "https://TEST-ACCOUNT.services.ai.azure.com/api/projects/test-project", + "https://test-account.SERVICES.ai.azure.com/api/projects/test-project", + "https://nested.test-account.services.ai.azure.com/api/projects/test-project", + "https://services.ai.azure.com/api/projects/test-project", + "https://test-account.services.ai.azure.com.evil.invalid/api/projects/test-project", + "https://test-account.azure.com/api/projects/test-project", + "https://user@test-account.services.ai.azure.com/api/projects/test-project", + "https://test-account.services.ai.azure.com/api/projects/test-project?", + "https://test-account.services.ai.azure.com/api/projects/test-project?version=8", + "https://test-account.services.ai.azure.com/api/projects/test-project#", + "https://test-account.services.ai.azure.com/api/projects/test-project#fragment", + "https://test-account.services.ai.azure.com/api/projects/%74est-project", + "https://test-account.services.ai.azure.com/api/projects/test%2Fproject", + "https://test-account.services.ai.azure.com/api/projects/../other", + "https://test-account.services.ai.azure.com/api/projects/..", + "https://test-account.services.ai.azure.com/api/projects/test-project/", + "https://test-account.services.ai.azure.com/api/projects/test-project/agents", + "https://test-account.services.ai.azure.com/api/projects/", + "https://test-account.services.ai.azure.com/API/projects/test-project", + "https://test-account.services.ai.azure.com/api/projects/test project", + "https://test-account.services.ai.azure.com/api/projects/über", + " https://test-account.services.ai.azure.com/api/projects/test-project", + "https://test-account.services.ai.azure.com/api/projects/test-project ", + } { + cfg := f.config + cfg.Target.ProjectEndpoint = endpoint + if validateHostedImageConfig(cfg) == nil { + t.Fatal("unsafe or noncanonical hosted endpoint accepted") + } + } + for _, version := range []string{"", "0", "01", "+1", "-1", "1.0", "1e1", "latest", "@latest", "1 ", "18446744073709551616"} { + cfg := f.config + cfg.Target.AgentVersion = version + if validateHostedImageConfig(cfg) == nil { + t.Fatal("noncanonical or unpinned hosted version accepted") + } + } + for name, mutate := range map[string]func(*hostedImageConfig){ + "protocol": func(c *hostedImageConfig) { c.Protocol = "other" }, + "zero deployment": func(c *hostedImageConfig) { c.DeploymentID = uuid.Nil.String() }, + "upper deployment": func(c *hostedImageConfig) { c.DeploymentID = "AAAAAAAA-AAAA-4AAA-8AAA-AAAAAAAAAAAA" }, + "short deployment": func(c *hostedImageConfig) { c.DeploymentID = strings.ReplaceAll(c.DeploymentID, "-", "") }, + "agent path": func(c *hostedImageConfig) { c.Target.AgentName = "test/agent" }, + "empty digest": func(c *hostedImageConfig) { c.AgentConfigurationDigest = "" }, + "uppercase digest": func(c *hostedImageConfig) { c.AgentConfigurationDigest = strings.ToUpper(c.AgentConfigurationDigest) }, + "zero public key": func(c *hostedImageConfig) { + c.SigningPublicKey = base64.RawURLEncoding.EncodeToString(make([]byte, ed25519.PublicKeySize)) + }, + "padded public key": func(c *hostedImageConfig) { c.SigningPublicKey += "=" }, + } { + t.Run(name, func(t *testing.T) { + cfg := f.config + mutate(&cfg) + if validateHostedImageConfig(cfg) == nil { + t.Fatal("invalid hosted image configuration accepted") + } + }) + } +} + +func TestHostedBootstrapAcceptsFrozenEnvironment(t *testing.T) { + f := newHostedProtocolFixture(t) + if validateHostedBootstrap(f.config, f.bootstrap) != nil { + t.Fatal("valid bootstrap rejected") + } + f.bootstrap.Environment["ORKA_ACP_WORKSPACE_INTENT"] = "write" + f.bootstrap.Environment["ORKA_ACP_MODEL_CONTEXT_LIMIT"] = "128000" + f.bootstrap.Environment["ORKA_ACP_MODEL_OUTPUT_LIMIT"] = "16384" + if validateHostedBootstrap(f.config, f.bootstrap) != nil { + t.Fatal("valid model limits and write workspace rejected") + } + for _, size := range []int{32, 16 << 10} { + for _, field := range []*string{&f.bootstrap.ControllerToken, &f.bootstrap.CapabilitySecret, &f.bootstrap.ProviderToken} { + *field = strings.Repeat("x", size) + } + if validateHostedBootstrap(f.config, f.bootstrap) != nil { + t.Fatal("valid credential size boundary rejected") + } + } +} + +func TestHostedBootstrapRejectsMissingAndAdditionalAuthority(t *testing.T) { + f := newHostedProtocolFixture(t) + for key := range f.bootstrap.Environment { + t.Run("missing "+key, func(t *testing.T) { + body := f.bootstrap + body.Environment = maps.Clone(body.Environment) + delete(body.Environment, key) + if validateHostedBootstrap(f.config, body) == nil { + t.Fatal("required bootstrap environment key omitted") + } + }) + } + for _, key := range []string{ + "HOME", "PATH", "LD_PRELOAD", "HTTP_PROXY", "AZURE_CLIENT_ID", "IDENTITY_ENDPOINT", "IDENTITY_HEADER", + "FOUNDRY_AGENT_SESSION_ID", "ORKA_ACP_WORKSPACE_ROOT", "ORKA_ACP_EXEC_HELPER", + "ORKA_ACP_PROCESS_COMMAND", "ORKA_ACP_LISTEN_ADDR", "ORKA_ACP_MCP_BROKER_URL", + "ORKA_ACP_CONTROLLER_TOKEN", "ORKA_ACP_PROVIDER_TOKEN", "ORKA_ACP_CAPABILITY_SECRET", + "ORKA_ACP_SUPERVISOR_BOOT_ID", "ORKA_ACP_RUNTIME_INSTANCE_ID", "ORKA_ACP_RUNTIME_SESSION_UID", + "ORKA_FOUNDRY_ACP_PROVIDER_BASE_URL", "ORKA_FOUNDRY_ACP_PROVIDER_TOKEN", "orka_acp_model", "UNKNOWN", + } { + t.Run("additional "+key, func(t *testing.T) { + body := f.bootstrap + body.Environment = maps.Clone(body.Environment) + body.Environment[key] = "synthetic-test-setting" + if validateHostedBootstrap(f.config, body) == nil { + t.Fatal("additional bootstrap environment authority accepted") + } + }) + } + f.bootstrap.Environment = nil + if validateHostedBootstrap(f.config, f.bootstrap) == nil { + t.Fatal("missing bootstrap environment accepted") + } +} + +func TestHostedBootstrapRejectsInvalidValues(t *testing.T) { + f := newHostedProtocolFixture(t) + for key, values := range map[string][]string{ + "ORKA_ACP_PROVIDER": {"", "codex", "Foundry", "foundry "}, + "ORKA_ACP_MODEL": {"", " ", " model", "model ", "model\n", string([]byte{0xff}), strings.Repeat("m", 513)}, + "ORKA_ACP_WORKSPACE_INTENT": {"", "Read", "read-write", "execute"}, + "ORKA_ACP_FOUNDRY_ADAPTER_DIGEST": {"", "invalid", "sha256:" + strings.Repeat("A", 64)}, + "ORKA_ACP_AGENT_CONFIGURATION_DIGEST": {"", brokerSHA([]byte("other agent configuration"))}, + "ORKA_ACP_TOOL_POLICY_DIGEST": {"", "invalid"}, + "ORKA_ACP_APPROVAL_POLICY_DIGEST": {"", "invalid"}, + "ORKA_ACP_MCP_CONFIGURATION_DIGEST": {"", "invalid"}, + "ORKA_ACP_PROXY_CREDENTIAL_ROLE": {"", "task-managed", "operator-managed "}, + "ORKA_ACP_PROXY_CREDENTIAL_SCOPE": {"", "task", "session"}, + "ORKA_ACP_RESOURCE_CLASS": {"", "privileged", "internal"}, + "ORKA_ACP_TRUST_NAMESPACE": {"", "Upper", "with.dot", "with_underscore", "-prefix", "suffix-", strings.Repeat("n", 64)}, + "ORKA_ACP_CONTROLLER_EPOCH": {"", "0", "01", "+1", "-1", "1.0", "1e1", "18446744073709551616"}, + "ORKA_ACP_RUNTIME_POOL_GENERATION": {"", "0", "01", "+1", "-1", "1.0", "1e1", "18446744073709551616"}, + "ORKA_ACP_RUNTIME_POOL_UID": {"", "pool", uuid.Nil.String(), "AAAAAAAA-AAAA-4AAA-8AAA-AAAAAAAAAAAA"}, + } { + t.Run(key, func(t *testing.T) { + for _, value := range values { + body := f.bootstrap + body.Environment = maps.Clone(body.Environment) + body.Environment[key] = value + if validateHostedBootstrap(f.config, body) == nil { + t.Fatal("invalid bootstrap environment value accepted") + } + } + }) + } + for _, limits := range [][2]string{ + {"1", ""}, {"", "1"}, {"0", "1"}, {"1", "0"}, {"01", "1"}, {"1", "01"}, + {"-1", "1"}, {"1", "+1"}, {"1.0", "1"}, {"1", "1e1"}, {"18446744073709551616", "1"}, + } { + body := f.bootstrap + body.Environment = maps.Clone(body.Environment) + if limits[0] != "" { + body.Environment["ORKA_ACP_MODEL_CONTEXT_LIMIT"] = limits[0] + } + if limits[1] != "" { + body.Environment["ORKA_ACP_MODEL_OUTPUT_LIMIT"] = limits[1] + } + if validateHostedBootstrap(f.config, body) == nil { + t.Fatal("invalid or partial model limits accepted") + } + } + badConfig := f.config + badConfig.AgentConfigurationDigest = "invalid" + if validateHostedBootstrap(badConfig, f.bootstrap) == nil { + t.Fatal("bootstrap validated against invalid image configuration") + } +} + +func TestHostedBootstrapRejectsMalformedCredentials(t *testing.T) { + f := newHostedProtocolFixture(t) + for name, set := range map[string]func(*hostedBootstrap, string){ + "controller": func(b *hostedBootstrap, value string) { b.ControllerToken = value }, + "capability": func(b *hostedBootstrap, value string) { b.CapabilitySecret = value }, + "provider": func(b *hostedBootstrap, value string) { b.ProviderToken = value }, + } { + t.Run(name, func(t *testing.T) { + for _, value := range []string{ + "", strings.Repeat("x", 31), strings.Repeat("x", (16<<10)+1), + strings.Repeat("x", 32) + " ", strings.Repeat("x", 32) + "\t", + strings.Repeat("x", 32) + "\n", strings.Repeat("x", 32) + "\x00", + strings.Repeat("x", 32) + "\u00a0", strings.Repeat("x", 32) + "\u0085", + strings.Repeat("x", 32) + string([]byte{0xff}), + } { + body := f.bootstrap + set(&body, value) + if validateHostedBootstrap(f.config, body) == nil { + t.Fatal("malformed bootstrap credential accepted") + } + } + }) + } +} diff --git a/hosted_proxy.go b/hosted_proxy.go new file mode 100644 index 0000000..a841d67 --- /dev/null +++ b/hosted_proxy.go @@ -0,0 +1,107 @@ +package main + +import ( + "crypto/tls" + "io" + "log" + "net" + "net/http" + "net/http/httputil" + "net/url" + "path" + "strings" + "time" +) + +// A fresh HTTP/1 connection for each local hop prevents the standard +// transport's stale-connection retry. The WebSocket hop uses ClientConn +// directly and likewise never retries an accepted request. +func newHostedLocalTransport() *http.Transport { + return &http.Transport{ + Proxy: nil, DialContext: (&net.Dialer{Timeout: 5 * time.Second}).DialContext, + DisableKeepAlives: true, DisableCompression: true, ForceAttemptHTTP2: false, + TLSHandshakeTimeout: 5 * time.Second, MaxResponseHeaderBytes: 64 << 10, + TLSNextProto: make(map[string]func(string, *tls.Conn) http.RoundTripper), + } +} + +func newHostedProxy(target string, transport http.RoundTripper, allow func(*http.Request) bool) (http.Handler, error) { + base, err := url.Parse(target) + if err != nil || base.Host == "" || base.Path != "" || (base.Scheme != "http" && base.Scheme != "https") { + return nil, errHostedInvalid + } + proxy := &httputil.ReverseProxy{ + Rewrite: func(request *httputil.ProxyRequest) { + request.Out.URL.Scheme = base.Scheme + request.Out.URL.Host = base.Host + request.Out.Host = base.Host + request.Out.GetBody = nil + // Preserve authorization exactly. In particular, the unsigned + // Foundry context is not authority to inject a broker credential. + }, + Transport: transport, FlushInterval: -1, + ErrorLog: log.New(io.Discard, "", 0), + ErrorHandler: func(w http.ResponseWriter, _ *http.Request, _ error) { + http.Error(w, "hosted transport unavailable; acceptance may be unknown", http.StatusBadGateway) + }, + } + return http.HandlerFunc(func(w http.ResponseWriter, r *http.Request) { + if !hostedRequestPathValid(r) || !allow(r) || r.Header.Get("Upgrade") != "" || r.Method == http.MethodConnect { + http.Error(w, "hosted route denied", http.StatusForbidden) + return + } + proxy.ServeHTTP(w, r) + }), nil +} + +func hostedRequestPathValid(r *http.Request) bool { + return r.URL != nil && r.URL.User == nil && r.URL.Opaque == "" && + r.URL.RawPath == "" && r.URL.Fragment == "" && len(r.URL.Path) <= 4096 && + len(r.URL.RawQuery) <= 8192 && path.Clean(r.URL.Path) == r.URL.Path && + strings.HasPrefix(r.URL.Path, "/") && !strings.ContainsAny(r.URL.Path, "\\\x00") +} + +func hostedV2Route(r *http.Request) bool { + if !strings.HasPrefix(r.URL.Path, "/v2/") { + return false + } + switch r.Method { + case http.MethodGet, http.MethodPost, http.MethodPut, http.MethodPatch, http.MethodDelete: + return true + default: + return false + } +} + +func hostedBrokerRoute(r *http.Request) bool { + if r.URL.RawQuery != "" { + return false + } + if r.URL.Path == brokerStatusPath { + return r.Method == http.MethodGet + } + if r.Method != http.MethodPost { + return false + } + switch r.URL.Path { + case brokerResponsesPath, brokerRenewPath, brokerSettlePath, brokerRetirePath: + return true + default: + return false + } +} + +func hostedOrkaRoute(r *http.Request) bool { + if r.URL.RawQuery != "" { + return false + } + if r.Method == http.MethodPost && (r.URL.Path == "/internal/v2/acp/mcp/tools/call" || + r.URL.Path == "/internal/v2/acp/artifact-authorizations") { + return true + } + const prefix = "/internal/v2/acp/artifacts/sha256/" + if (r.Method == http.MethodGet || r.Method == http.MethodPut || r.Method == http.MethodHead) && strings.HasPrefix(r.URL.Path, prefix) { + return brokerDigestValid("sha256:" + strings.TrimPrefix(r.URL.Path, prefix)) + } + return false +} diff --git a/hosted_remote.go b/hosted_remote.go new file mode 100644 index 0000000..839ade4 --- /dev/null +++ b/hosted_remote.go @@ -0,0 +1,269 @@ +package main + +import ( + "bytes" + "context" + "encoding/base64" + "encoding/json" + "io" + "net/http" + "net/url" + "strings" + "time" + + "github.com/gorilla/websocket" + "golang.org/x/net/http/httpguts" +) + +// Azure authorizes the exact hosted session. The in-band signature separately +// authorizes the supervisor bootstrap; Foundry strips ordinary custom headers. +func (g *hostedGateway) accessToken(ctx context.Context) (string, error) { + token, err := g.provider.AccessToken(ctx) + if err != nil || !httpguts.ValidHeaderFieldValue(token) { + return "", errHostedInvalid + } + parts := strings.Split(token, ".") + if len(parts) != 3 { + return "", errHostedInvalid + } + raw, err := base64.RawURLEncoding.DecodeString(parts[1]) + defer clear(raw) + var claims struct { + Audience string `json:"aud"` + Tenant string `json:"tid"` + Object string `json:"oid"` + App string `json:"appid"` + AZP string `json:"azp"` + } + if err != nil || acpDecodeStruct(raw, &claims, false) != nil || + strings.TrimRight(claims.Audience, "/") != "https://ai.azure.com" || + !hostedUUIDValid(claims.Tenant) || !hostedUUIDValid(claims.Object) { + return "", errHostedInvalid + } + digest := brokerJSONDigest(claims) + g.mu.Lock() + defer g.mu.Unlock() + if g.ledger.PrincipalDigest == "" { + next := g.ledger + next.PrincipalDigest = digest + if g.store.save(next) != nil { + return "", errHostedInvalid + } + g.ledger = next + } + if g.ledger.PrincipalDigest != digest { + return "", errHostedInvalid + } + return token, nil +} + +func (g *hostedGateway) remoteURL(suffix string) string { + target := g.cfg.Image.Target + return target.ProjectEndpoint + "/agents/" + url.PathEscape(target.AgentName) + suffix + "?api-version=v1" +} + +func (g *hostedGateway) remoteJSON(ctx context.Context, method, suffix string, body []byte, target any) (int, error) { + request, err := g.prepareRemoteJSON(ctx, method, suffix, body) + if err != nil { + return 0, err + } + return g.sendRemoteJSON(request, target) +} + +func (g *hostedGateway) prepareRemoteJSON(ctx context.Context, method, suffix string, body []byte) (*http.Request, error) { + if ctx.Err() != nil { + return nil, errHostedInvalid + } + token, err := g.accessToken(ctx) + if err != nil { + return nil, err + } + request, err := http.NewRequestWithContext(ctx, method, g.remoteURL(suffix), bytes.NewReader(body)) + if err != nil { + return nil, errHostedInvalid + } + request.GetBody = nil + request.Header.Set("Authorization", "Bearer "+token) + request.Header.Set("Foundry-Features", "HostedAgents=V1Preview") + request.Header.Set("Content-Type", "application/json") + request.Header.Set("Accept", "application/json") + if ctx.Err() != nil { + return nil, errHostedInvalid + } + return request, nil +} + +func (g *hostedGateway) sendRemoteJSON(request *http.Request, target any) (int, error) { + response, err := g.httpClient.Do(request) + if err != nil { + return 0, errHostedInvalid + } + defer response.Body.Close() //nolint:errcheck + data, err := io.ReadAll(io.LimitReader(response.Body, acpMaxConfigBytes+1)) + defer clear(data) + if err != nil || len(data) > acpMaxConfigBytes { + return response.StatusCode, errHostedInvalid + } + if target != nil && response.StatusCode >= 200 && response.StatusCode < 300 && acpDecodeStruct(data, target, false) != nil { + return response.StatusCode, errHostedInvalid + } + return response.StatusCode, nil +} + +func (g *hostedGateway) validateRemote(ctx context.Context) error { + var agent struct { + Name string `json:"name"` + Endpoint struct { + Schemes []json.RawMessage `json:"authorization_schemes"` + } `json:"agent_endpoint"` + } + status, err := g.remoteJSON(ctx, http.MethodGet, "", nil, &agent) + if err != nil || status != http.StatusOK || agent.Name != g.cfg.Image.Target.AgentName || len(agent.Endpoint.Schemes) == 0 { + return errHostedInvalid + } + for _, scheme := range agent.Endpoint.Schemes { + var value struct { + Type string `json:"type"` + } + if acpDecodeStruct(scheme, &value, true) != nil || !strings.EqualFold(value.Type, "entra") { + return errHostedInvalid + } + } + var version struct { + Name string `json:"name"` + Version string `json:"version"` + Status string `json:"status"` + Definition struct { + Kind string `json:"kind"` + Container struct { + Image string `json:"image"` + } `json:"container_configuration"` + Protocols []struct { + Protocol string `json:"protocol"` + Version string `json:"version"` + } `json:"protocol_versions"` + } `json:"definition"` + } + status, err = g.remoteJSON(ctx, http.MethodGet, "/versions/"+g.cfg.Image.Target.AgentVersion, nil, &version) + if err != nil || status != http.StatusOK || version.Name != g.cfg.Image.Target.AgentName || + version.Version != g.cfg.Image.Target.AgentVersion || version.Status != "active" || + version.Definition.Kind != "hosted" || version.Definition.Container.Image != g.cfg.ContainerImage { + return errHostedInvalid + } + ws := 0 + for _, protocol := range version.Definition.Protocols { + if protocol.Protocol == "invocations_ws" { + if protocol.Version != "2.0.0" { + return errHostedInvalid + } + ws++ + } + } + if ws != 1 { + return errHostedInvalid + } + return nil +} + +func (g *hostedGateway) ensureSession(ctx context.Context) error { + // An uncertain create is never replayed or adopted from a subsequent GET. + if g.ledger.CreateAttempted && !g.ledger.SessionCreated { + return errHostedInvalid + } + var session brokerRemoteSession + status, err := g.remoteJSON(ctx, http.MethodGet, brokerSessionSuffix(g.cfg.SessionID), nil, &session) + if err != nil { + return err + } + if !g.ledger.CreateAttempted { + if status != http.StatusNotFound { + return errHostedInvalid + } + body, _ := json.Marshal(map[string]any{"agent_session_id": g.cfg.SessionID, + "version_indicator": map[string]string{"type": "version_ref", "agent_version": g.cfg.Image.Target.AgentVersion}}) + request, prepareErr := g.prepareRemoteJSON(ctx, http.MethodPost, "/endpoint/sessions", body) + if prepareErr != nil || ctx.Err() != nil { + return errHostedInvalid + } + // Authentication and local request preparation cannot strand an + // unsent creation. Once this intent is durable, submission errors + // remain ambiguous and never authorize retry or adoption. + next := g.ledger + next.CreateAttempted = true + if g.store.save(next) != nil { + return errHostedInvalid + } + g.ledger = next + status, err = g.sendRemoteJSON(request, &session) + if err == nil && brokerDefiniteRejection(status) { + // A complete admission rejection proves that this attempt created no + // session. Persist that result before permitting a later startup. + // Transport errors and incomplete responses retain the original intent. + next = g.ledger + next.CreateAttempted = false + if g.store.save(next) != nil { + return errHostedInvalid + } + g.ledger = next + return errHostedInvalid + } + if err != nil || status != http.StatusCreated || !g.sessionMatches(session) { + return errHostedInvalid + } + next = g.ledger + next.SessionCreated = true + if g.store.save(next) != nil { + return errHostedInvalid + } + g.ledger = next + status, err = g.remoteJSON(ctx, http.MethodGet, brokerSessionSuffix(g.cfg.SessionID), nil, &session) + } + if err != nil || status != http.StatusOK || !g.sessionMatches(session) || session.Status != "active" { + return errHostedInvalid + } + return nil +} + +func (g *hostedGateway) sessionMatches(value brokerRemoteSession) bool { + return value.ID == g.cfg.SessionID && value.Version.Type == "version_ref" && value.Version.Version == g.cfg.Image.Target.AgentVersion +} + +func (g *hostedGateway) openChannel(ctx context.Context, role string) (*websocket.Conn, hostedChallenge, error) { + var challenge hostedChallenge + token, err := g.accessToken(ctx) + if err != nil { + return nil, challenge, err + } + endpoint := strings.Replace(g.remoteURL("/endpoint/protocols/invocations_ws"), "https://", "wss://", 1) + + "&agent_session_id=" + url.QueryEscape(g.cfg.SessionID) + headers := http.Header{"Authorization": []string{"Bearer " + token}, "Foundry-Features": []string{"HostedAgents=V1Preview"}} + ws, response, err := g.dial(ctx, endpoint, headers) + if err != nil { + if response != nil && response.Body != nil { + _ = response.Body.Close() + } + return nil, challenge, errHostedInvalid + } + observation := observeHostedChannel(ws, role, time.Now(), g.channelLog) + if role == "forward" { + g.forwardObservation = observation + } else { + g.reverseObservation = observation + } + ws.SetReadLimit(hostedMaxHandshakeBytes) + _ = ws.SetReadDeadline(time.Now().Add(30 * time.Second)) + if readHostedWSJSON(ws, &challenge) != nil || challenge.Protocol != hostedProtocol || + challenge.DeploymentID != g.cfg.Image.DeploymentID || challenge.ConfigurationDigest != brokerJSONDigest(g.cfg.Image) || + challenge.AgentName != g.cfg.Image.Target.AgentName || challenge.AgentVersion != g.cfg.Image.Target.AgentVersion || + challenge.SessionID != g.cfg.SessionID || !hostedUUIDValid(challenge.BootID) { + _ = ws.Close() + observation.finish() + return nil, hostedChallenge{}, errHostedInvalid + } + if _, ok := hostedCanonicalBytes(challenge.Nonce, 32); !ok { + _ = ws.Close() + observation.finish() + return nil, hostedChallenge{}, errHostedInvalid + } + return ws, challenge, nil +} diff --git a/hosted_review_pairing_test.go b/hosted_review_pairing_test.go new file mode 100644 index 0000000..9a74377 --- /dev/null +++ b/hosted_review_pairing_test.go @@ -0,0 +1,159 @@ +package main + +import ( + "context" + "net/http" + "testing" + "testing/synctest" + "time" + + "github.com/google/uuid" + "github.com/gorilla/websocket" +) + +func TestHostedReviewServerAllowsHelloAfterOtherChannelSetup(t *testing.T) { + f := newHostedServerTestFixture(t, false) + synctest.Test(t, func(t *testing.T) { + f.server.cancel() + f.server.ctx, f.server.cancel = context.WithCancel(t.Context()) + defer f.server.close() //nolint:errcheck + pairID, digest := uuid.NewString(), f.protocol.hello.BootstrapDigest + forward := hostedReviewPipeConnect(t, f) + // The gateway verifies both challenges before signing. A second dial + // may take up to 120 seconds, exceeding the old first-hello window. + time.Sleep(31 * time.Second) + reverse := hostedReviewPipeConnect(t, f) + if forward.WriteJSON(f.hello(t, "forward", pairID, digest)) != nil { + t.Fatal("valid first hello could not be delivered after slow second-channel setup") + } + hostedReviewReadAck(t, f, forward, "forward", pairID) + if reverse.WriteJSON(f.hello(t, "reverse", pairID, digest)) != nil { + t.Fatal("valid second hello could not be delivered") + } + hostedReviewReadAck(t, f, reverse, "reverse", pairID) + synctest.Wait() + if f.server.ctx.Err() != nil || f.calls.Load() != 0 || !f.state(func(p *hostedPair) bool { + return p != nil && p.forwardAcked && p.reverseAcked && p.bootstrap == nil && !p.running + }) { + t.Fatal("slow channel pairing expired or launched before bootstrap") + } + }) +} + +func hostedReviewPipeConnect(t *testing.T, f *hostedServerTestFixture) *websocket.Conn { + t.Helper() + ws, response, err := hostedHandshakePipeDial(t, t.Context(), f.server, "ws://fixture.invalid/invocations_ws", nil) + if response != nil && response.Body != nil { + _ = response.Body.Close() + } + if err != nil { + t.Fatal("could not establish the in-memory hosted channel") + } + t.Cleanup(func() { _ = ws.Close() }) + _ = ws.SetReadDeadline(time.Now().Add(5 * time.Minute)) + var challenge hostedChallenge + if readHostedWSJSON(ws, &challenge) != nil || challenge != f.server.challenge { + t.Fatal("server did not send the exact in-memory challenge") + } + return ws +} + +func hostedReviewReadAck(t *testing.T, f *hostedServerTestFixture, ws *websocket.Conn, role, pairID string) { + t.Helper() + var ack hostedAccepted + if readHostedWSJSON(ws, &ack) != nil || ack != (hostedAccepted{Protocol: hostedProtocol, + PairID: pairID, Role: role, BootID: f.server.challenge.BootID, BootstrapDigest: f.protocol.hello.BootstrapDigest}) { + t.Fatal("valid authenticated role was not acknowledged") + } +} + +func TestHostedReviewServerAllowsBootstrapAfterSlowOtherRole(t *testing.T) { + for _, test := range []struct { + name string + hello, ack time.Duration + }{ + {name: "reverse hello near its deadline", hello: 29 * time.Second}, + {name: "reverse acknowledgment near its deadline", hello: 22 * time.Second, ack: 9 * time.Second}, + } { + t.Run(test.name, func(t *testing.T) { + f := newHostedServerTestFixture(t, false) + synctest.Test(t, func(t *testing.T) { + f.server.cancel() + f.server.ctx, f.server.cancel = context.WithCancel(t.Context()) + defer f.server.close() //nolint:errcheck + pairID, digest := uuid.NewString(), f.protocol.hello.BootstrapDigest + forward := hostedReviewPipeConnect(t, f) + if forward.WriteJSON(f.hello(t, "forward", pairID, digest)) != nil { + t.Fatal("could not send the first authenticated role") + } + hostedReviewReadAck(t, f, forward, "forward", pairID) + time.Sleep(2 * time.Second) + reverse := hostedReviewPipeConnect(t, f) + time.Sleep(test.hello) + if reverse.WriteJSON(f.hello(t, "reverse", pairID, digest)) != nil { + t.Fatal("reverse hello inside its own window could not be delivered") + } + time.Sleep(test.ack) + hostedReviewReadAck(t, f, reverse, "reverse", pairID) + reverseConn := newHostedWSConn(reverse) + go serveHostedHTTP2(f.server.ctx, reverseConn, http.HandlerFunc(func(w http.ResponseWriter, _ *http.Request) { + w.WriteHeader(http.StatusNoContent) + })) + f.sendBootstrap(t, forward) + <-f.server.ctx.Done() // The deliberate fixture runner failure ends this lifetime. + f.server.mu.Lock() + pair := f.server.pair + f.server.mu.Unlock() + <-pair.stopped + if f.calls.Load() != 1 || !f.state(func(p *hostedPair) bool { return p.running && p.bootstrap != nil }) { + t.Fatal("valid slow pairing failed to deliver exactly one bootstrap to the runner") + } + }) + }) + } +} + +func TestHostedReviewAuthenticatedPairingHasOneBound(t *testing.T) { + for _, role := range []string{"forward", "reverse"} { + t.Run(role, func(t *testing.T) { + f := newHostedServerTestFixture(t, false) + synctest.Test(t, func(t *testing.T) { + f.server.cancel() + f.server.ctx, f.server.cancel = context.WithCancel(t.Context()) + defer f.server.close() //nolint:errcheck + pairID := uuid.NewString() + ws := hostedReviewPipeConnect(t, f) + if ws.WriteJSON(f.hello(t, role, pairID, f.protocol.hello.BootstrapDigest)) != nil { + t.Fatal("could not send the authenticated role") + } + hostedReviewReadAck(t, f, ws, role, pairID) + synctest.Wait() + time.Sleep(4*time.Minute - time.Second) + if f.server.ctx.Err() != nil || f.calls.Load() != 0 { + t.Fatal("incomplete authenticated pair expired before its setup window or launched early") + } + time.Sleep(time.Second) + synctest.Wait() + if f.server.ctx.Err() == nil || f.calls.Load() != 0 { + t.Fatal("incomplete pair outlived its single setup window or launched a supervisor") + } + }) + }) + } +} + +func TestHostedReviewUnauthenticatedPairingWaitIsBounded(t *testing.T) { + f := newHostedServerTestFixture(t, false) + synctest.Test(t, func(t *testing.T) { + f.server.cancel() + f.server.ctx, f.server.cancel = context.WithCancel(t.Context()) + defer f.server.close() //nolint:errcheck + ws := hostedReviewPipeConnect(t, f) + time.Sleep(4 * time.Minute) + synctest.Wait() + if readHostedWSJSON(ws, &hostedAccepted{}) == nil || f.server.ctx.Err() != nil || f.calls.Load() != 0 || + !f.state(func(p *hostedPair) bool { return p == nil }) { + t.Fatal("expired unauthenticated channel stayed open or consumed authenticated ownership") + } + }) +} diff --git a/hosted_review_shutdown_test.go b/hosted_review_shutdown_test.go new file mode 100644 index 0000000..7c06957 --- /dev/null +++ b/hosted_review_shutdown_test.go @@ -0,0 +1,215 @@ +package main + +import ( + "context" + "errors" + "maps" + "net" + "net/http" + "sync" + "testing" + "time" + + "github.com/google/uuid" +) + +func TestHostedReviewLocalStopPreservesCleanupRelays(t *testing.T) { + hostedServerTestHealth(t, http.StatusOK) + f := newHostedServerTestFixture(t, true) + stopRequested, releaseStop := make(chan struct{}), make(chan struct{}) + cleanupOK := make(chan bool, 1) + release := sync.OnceFunc(func() { close(releaseStop) }) + defer release() + f.server.runner = func(ctx context.Context, environment map[string]string) (<-chan error, error) { + f.calls.Add(1) + f.captured <- maps.Clone(environment) + done := make(chan error, 1) + go func() { + <-ctx.Done() + close(stopRequested) + <-releaseStop + transport := newHostedLocalTransport() + defer transport.CloseIdleConnections() + client := &http.Client{Transport: transport, Timeout: 2 * time.Second} + ok := true + for _, target := range []string{ + "http://" + hostedBrokerRelayAddr + brokerSettlePath, + "http://" + hostedOrkaRelayAddr + "/internal/v2/acp/artifact-authorizations", + } { + request, _ := http.NewRequest(http.MethodPost, target, nil) + response, err := client.Do(request) + if err != nil { + ok = false + continue + } + ok = ok && response.StatusCode == http.StatusNoContent + _ = response.Body.Close() + } + cleanupOK <- ok + if ok { + done <- nil + } else { + done <- errHostedInvalid + } + close(done) + }() + return done, nil + } + forward, reverse, client := hostedServerTestRunningPair(t, f) + returned := make(chan error, 1) + go func() { returned <- serveHostedHTTP(f.server.ctx, "127.0.0.1:0", f.server) }() + f.server.cancel() + hostedTestDone(t, stopRequested) + select { + case <-returned: + t.Fatal("entry point did not join the supervisor's pending cleanup") + case <-time.After(25 * time.Millisecond): + } + ctx, cancel := context.WithTimeout(t.Context(), 2*time.Second) + defer cancel() + request, _ := http.NewRequestWithContext(ctx, http.MethodGet, "http://supervisor/v2/health", nil) + response, err := client.RoundTrip(request) + if err != nil { + t.Error("local shutdown destroyed the healthy forward transport before cleanup") + } else { + _ = response.Body.Close() + if response.StatusCode != http.StatusServiceUnavailable { + t.Error("stopping lifetime admitted a new forward request") + } + } + release() + select { + case ok := <-cleanupOK: + if !ok { + t.Error("ordinary local shutdown destroyed a healthy relay needed by supervisor cleanup") + } + case <-time.After(5 * time.Second): + t.Fatal("bounded fixture cleanup did not return") + } + select { + case err := <-returned: + if err != nil { + t.Error("clean joined supervisor shutdown was rejected") + } + case <-time.After(3 * time.Second): + t.Fatal("entry point did not finish after supervisor cleanup") + } + hostedTestDone(t, forward.Done()) + hostedTestDone(t, reverse.Done()) + if f.calls.Load() != 1 { + t.Fatal("local shutdown relaunched its supervisor") + } +} + +func TestHostedReviewLocalRelayStartupFailureIsNotCleanExit(t *testing.T) { + for _, address := range []string{hostedBrokerRelayAddr, hostedOrkaRelayAddr} { + name := "broker relay" + if address == hostedOrkaRelayAddr { + name = "Orka relay" + } + t.Run(name, func(t *testing.T) { + reserved, err := net.Listen("tcp", address) + if err != nil { + t.Fatal("could not reserve the local fixture relay port") + } + defer reserved.Close() //nolint:errcheck + f := newHostedServerTestFixture(t, false) + pairID, digest := uuid.NewString(), brokerJSONDigest(f.protocol.bootstrap) + forward := f.claim(t, "forward", pairID, digest) + _ = hostedServerReverse(t, f.claim(t, "reverse", pairID, digest)) + f.sendBootstrap(t, forward) + hostedTestDone(t, f.server.ctx.Done()) + f.server.mu.Lock() + pair := f.server.pair + f.server.mu.Unlock() + hostedTestDone(t, pair.stopped) + if !errors.Is(f.server.close(), errHostedInvalid) { + t.Error("local relay bind failure was reported as a clean hosted exit") + } + if f.calls.Load() != 0 { + t.Fatal("failed local startup launched or retried a supervisor") + } + }) + } +} + +func TestHostedReviewPeerCancellationWhileStartingIsNotLocalFailure(t *testing.T) { + for _, role := range []string{"forward", "reverse"} { + t.Run(role, func(t *testing.T) { + healthCalled := hostedServerTestHealth(t, http.StatusServiceUnavailable) + f := newHostedServerTestFixture(t, true) + f.server.runner = func(ctx context.Context, environment map[string]string) (<-chan error, error) { + f.calls.Add(1) + f.captured <- maps.Clone(environment) + done := make(chan error, 1) + go func() { <-ctx.Done(); done <- nil; close(done) }() + return done, nil + } + pairID, digest := uuid.NewString(), brokerJSONDigest(f.protocol.bootstrap) + forward := f.claim(t, "forward", pairID, digest) + reverse := hostedServerReverse(t, f.claim(t, "reverse", pairID, digest)) + f.sendBootstrap(t, forward) + _ = f.receiveRunner(t) + hostedTestDone(t, healthCalled) + if role == "forward" { + hostedTestGoingAway(t, forward) + } else { + hostedTestGoingAway(t, reverse.ws) + } + hostedTestDone(t, f.server.ctx.Done()) + if f.server.close() != nil || f.calls.Load() != 1 { + t.Fatal("peer-canceled startup was classified as local failure or relaunched") + } + }) + } +} + +func TestHostedReviewChannelLossDuringCleanupClosesPair(t *testing.T) { + for _, role := range []string{"forward", "reverse"} { + t.Run(role, func(t *testing.T) { + hostedServerTestHealth(t, http.StatusOK) + f := newHostedServerTestFixture(t, true) + stopRequested, releaseStop := make(chan struct{}), make(chan struct{}) + release := sync.OnceFunc(func() { close(releaseStop) }) + defer release() + f.server.runner = func(ctx context.Context, environment map[string]string) (<-chan error, error) { + f.calls.Add(1) + f.captured <- maps.Clone(environment) + done := make(chan error, 1) + go func() { + <-ctx.Done() + close(stopRequested) + <-releaseStop + done <- nil + close(done) + }() + return done, nil + } + forward, reverse, _ := hostedServerTestRunningPair(t, f) + returned := make(chan error, 1) + go func() { returned <- f.server.close() }() + hostedTestDone(t, stopRequested) + channel := forward + if role == "reverse" { + channel = reverse + } + hostedTestGoingAway(t, channel.ws) + hostedTestDone(t, forward.Done()) + hostedTestDone(t, reverse.Done()) + select { + case <-returned: + t.Error("channel loss returned before joining pending supervisor cleanup") + default: + } + release() + select { + case err := <-returned: + if err != nil || f.calls.Load() != 1 { + t.Fatal("channel loss during local cleanup failed to join or replayed the supervisor") + } + case <-time.After(3 * time.Second): + t.Fatal("channel loss during cleanup left the lifetime pending") + } + }) + } +} diff --git a/hosted_server.go b/hosted_server.go new file mode 100644 index 0000000..c83c2ea --- /dev/null +++ b/hosted_server.go @@ -0,0 +1,523 @@ +package main + +import ( + "bufio" + "context" + "crypto/rand" + "encoding/base64" + "encoding/json" + "io" + "log" + "net" + "net/http" + "net/url" + "os" + "strings" + "sync" + "sync/atomic" + "time" + + "github.com/google/uuid" + "github.com/gorilla/websocket" +) + +type hostedReady struct { + Protocol string `json:"protocol"` + PairID string `json:"pairID"` + BootID string `json:"bootID"` + Ready bool `json:"ready"` +} + +type hostedSupervisorRunner func(context.Context, map[string]string) (<-chan error, error) + +const ( + // Both challenges must match before either role is signed. Allow the + // gateway's bounded initialization budget for the other channel's dial + // and handshake, then bound authenticated pairing by one shared deadline. + hostedInitializationWait = 4 * time.Minute + + // Orka allows 45 seconds for HTTP shutdown, then 45 seconds for session + // cleanup. Give that owner time to reap its separate ACP UID scopes. + hostedSupervisorStopGrace = 95 * time.Second + hostedSupervisorKillWait = 5 * time.Second + hostedSupervisorShutdownWait = hostedSupervisorStopGrace + hostedSupervisorKillWait + 5*time.Second +) + +type hostedServer struct { + cfg hostedImageConfig + challenge hostedChallenge + adapterDigest string + agentConfig []byte + runner hostedSupervisorRunner + ctx context.Context + cancel context.CancelFunc + mu sync.Mutex + pair *hostedPair + channelLog *log.Logger +} + +type hostedPair struct { + server *hostedServer + id, bootstrapDigest string + forward, reverse *websocket.Conn + forwardObservation, reverseObservation *hostedChannelObservation + forwardAcked, reverseAcked bool + setupDeadline time.Time + bootstrap *hostedBootstrap + started chan struct{} + done chan struct{} + stopped chan struct{} + shutdownErr error + once sync.Once + running bool +} + +func newHostedServer(ctx context.Context, cfg hostedImageConfig, getenv func(string) string, runner hostedSupervisorRunner) (*hostedServer, error) { + if validateHostedImageConfig(cfg) != nil || os.Geteuid() != 0 || + strings.TrimRight(getenv("FOUNDRY_PROJECT_ENDPOINT"), "/") != cfg.Target.ProjectEndpoint || + getenv("FOUNDRY_AGENT_NAME") != cfg.Target.AgentName || getenv("FOUNDRY_AGENT_VERSION") != cfg.Target.AgentVersion { + return nil, errHostedInvalid + } + sid, err := uuid.Parse(getenv("FOUNDRY_AGENT_SESSION_ID")) + if err != nil || !hostedUUIDValid(getenv("FOUNDRY_AGENT_SESSION_ID")) { + return nil, errHostedInvalid + } + digest := getenv("ORKA_ACP_FOUNDRY_ADAPTER_DIGEST") + agent, err := readHostedFile(acpConfigPath, acpMaxConfigBytes) + if err != nil || brokerSHA(agent) != cfg.AgentConfigurationDigest || !brokerDigestValid(digest) { + return nil, errHostedInvalid + } + var nonce [32]byte + if _, err := rand.Read(nonce[:]); err != nil { + return nil, errHostedInvalid + } + lifetime, cancel := context.WithCancel(ctx) + return &hostedServer{cfg: cfg, adapterDigest: digest, agentConfig: agent, runner: runner, + ctx: lifetime, cancel: cancel, challenge: hostedChallenge{Protocol: hostedProtocol, + DeploymentID: cfg.DeploymentID, ConfigurationDigest: brokerJSONDigest(cfg), + AgentName: cfg.Target.AgentName, AgentVersion: cfg.Target.AgentVersion, + SessionID: sid.String(), BootID: uuid.NewString(), Nonce: base64.RawURLEncoding.EncodeToString(nonce[:])}}, nil +} + +func (s *hostedServer) close() error { + s.cancel() + s.mu.Lock() + pair := s.pair + running := pair != nil && pair.running + s.mu.Unlock() + if pair == nil { + return nil + } + if !running { + pair.close() + return nil + } + timer := time.NewTimer(hostedSupervisorShutdownWait + time.Second) + defer timer.Stop() + select { + case <-pair.stopped: + s.mu.Lock() + defer s.mu.Unlock() + return pair.shutdownErr + case <-timer.C: + pair.close() + return errHostedInvalid + } +} + +func (s *hostedServer) ServeHTTP(w http.ResponseWriter, r *http.Request) { + if r.URL.Path == "/readiness" && r.Method == http.MethodGet && r.URL.RawQuery == "" { + if s.ctx.Err() != nil { + http.Error(w, "hosted lifetime closed", http.StatusServiceUnavailable) + return + } + w.Header().Set("Content-Type", "application/json") + _, _ = io.WriteString(w, `{"ready":true}`) + return + } + if r.URL.Path != "/invocations_ws" || r.Method != http.MethodGet || r.URL.RawPath != "" || !hostedRoutingQuery(r) { + http.NotFound(w, r) + return + } + if s.ctx.Err() != nil { + http.Error(w, "hosted lifetime closed", http.StatusServiceUnavailable) + return + } + upgrader := websocket.Upgrader{HandshakeTimeout: 10 * time.Second, + ReadBufferSize: 4096, WriteBufferSize: 4096, + CheckOrigin: func(request *http.Request) bool { return request.Header.Get("Origin") == "" }} + ws, err := upgrader.Upgrade(w, r, nil) + if err != nil { + return + } + openedAt := time.Now() + defer ws.Close() //nolint:errcheck + ws.SetReadLimit(hostedMaxHandshakeBytes) + _ = ws.SetWriteDeadline(time.Now().Add(10 * time.Second)) + if ws.WriteJSON(s.challenge) != nil { + return + } + var hello hostedHello + _ = ws.SetReadDeadline(time.Now().Add(hostedInitializationWait)) + if readHostedWSJSON(ws, &hello) != nil || verifyHostedHello(hello, s.challenge, s.cfg.SigningPublicKey, time.Now()) != nil { + _ = ws.WriteControl(websocket.CloseMessage, websocket.FormatCloseMessage(websocket.ClosePolicyViolation, "hosted authentication failed"), time.Now().Add(time.Second)) + return + } + observation := observeHostedChannel(ws, hello.Role, openedAt, s.channelLog) + defer observation.finish() + pair, err := s.claim(hello, ws, observation) + if err != nil { + _ = ws.WriteControl(websocket.CloseMessage, websocket.FormatCloseMessage(websocket.ClosePolicyViolation, "hosted pair conflicts"), time.Now().Add(time.Second)) + return + } + ack := hostedAccepted{Protocol: hostedProtocol, PairID: pair.id, Role: hello.Role, + BootID: s.challenge.BootID, BootstrapDigest: pair.bootstrapDigest} + _ = ws.SetWriteDeadline(time.Now().Add(10 * time.Second)) + if ws.WriteJSON(ack) != nil { + pair.close() + return + } + s.mu.Lock() + if hello.Role == "forward" { + pair.forwardAcked = true + } else { + pair.reverseAcked = true + } + s.mu.Unlock() + if hello.Role == "forward" { + _ = ws.SetReadDeadline(pair.setupDeadline) + kind, data, err := ws.ReadMessage() + var bootstrap hostedBootstrap + if err != nil || kind != websocket.TextMessage || brokerSHA(data) != pair.bootstrapDigest || + acpDecode(data, &bootstrap, true) != nil || validateHostedBootstrap(s.cfg, bootstrap) != nil || + bootstrap.Environment["ORKA_ACP_FOUNDRY_ADAPTER_DIGEST"] != s.adapterDigest { + clear(data) + pair.close() + return + } + clear(data) + if _, err := decodeACPAgentConfiguration(s.agentConfig, s.cfg.AgentConfigurationDigest, bootstrap.Environment["ORKA_ACP_MODEL"]); err != nil { + pair.close() + return + } + s.mu.Lock() + pair.bootstrap = &bootstrap + s.mu.Unlock() + } + pair.maybeStart() + timer := time.NewTimer(time.Until(pair.setupDeadline)) + defer timer.Stop() + select { + case <-pair.started: + case <-timer.C: + pair.close() + case <-s.ctx.Done(): + // A running supervisor owns bounded shutdown. Keep its healthy + // transports available until cleanup completes. + s.mu.Lock() + running := pair.running + s.mu.Unlock() + if !running { + pair.close() + } + } + <-pair.done +} + +func hostedRoutingQuery(r *http.Request) bool { + if len(r.URL.RawQuery) > 2048 { + return false + } + query, err := url.ParseQuery(r.URL.RawQuery) + if err != nil { + return false + } + for key, values := range query { + if (key != "api-version" && key != "agent_session_id") || len(values) != 1 || len(values[0]) > 512 || !acpSafeString(values[0], 512) { + return false + } + } + return true +} + +func readHostedWSJSON(ws *websocket.Conn, value any) error { + kind, data, err := ws.ReadMessage() + if err != nil || kind != websocket.TextMessage || len(data) > hostedMaxHandshakeBytes || acpDecode(data, value, true) != nil { + return errHostedInvalid + } + return nil +} + +func (s *hostedServer) claim(hello hostedHello, ws *websocket.Conn, observation *hostedChannelObservation) (*hostedPair, error) { + s.mu.Lock() + defer s.mu.Unlock() + if s.ctx.Err() != nil { + return nil, errHostedInvalid + } + if s.pair == nil { + s.pair = &hostedPair{server: s, id: hello.PairID, bootstrapDigest: hello.BootstrapDigest, + setupDeadline: time.Now().Add(hostedInitializationWait), + started: make(chan struct{}), done: make(chan struct{}), stopped: make(chan struct{})} + } + pair := s.pair + if pair.id != hello.PairID || pair.bootstrapDigest != hello.BootstrapDigest || pair.running { + return nil, errHostedInvalid + } + switch hello.Role { + case "forward": + if pair.forward != nil { + return nil, errHostedInvalid + } + pair.forward = ws + pair.forwardObservation = observation + case "reverse": + if pair.reverse != nil { + return nil, errHostedInvalid + } + pair.reverse = ws + pair.reverseObservation = observation + default: + return nil, errHostedInvalid + } + return pair, nil +} + +func (p *hostedPair) maybeStart() { + p.server.mu.Lock() + defer p.server.mu.Unlock() + if !p.running && p.forwardAcked && p.reverseAcked && p.bootstrap != nil && + p.server.ctx.Err() == nil && time.Now().Before(p.setupDeadline) { + p.running = true + close(p.started) + go p.run() + } +} + +func (p *hostedPair) close() { + p.once.Do(func() { + p.server.cancel() + p.server.mu.Lock() + forward, reverse := p.forward, p.reverse + forwardObservation, reverseObservation := p.forwardObservation, p.reverseObservation + p.server.mu.Unlock() + if forward != nil { + _ = forward.Close() + } + if reverse != nil { + _ = reverse.Close() + } + forwardObservation.finish() + reverseObservation.finish() + close(p.done) + }) +} + +func (p *hostedPair) run() { + defer close(p.stopped) + defer p.close() + s := p.server + // Local shutdown stops admission and the supervisor first. Its cleanup + // still needs the reverse relay, and the gateway requires both channels. + // Actual channel loss continues to close the pair immediately. + transportContext, stopTransport := context.WithCancel(context.WithoutCancel(s.ctx)) + defer stopTransport() + _ = p.forward.SetReadDeadline(time.Time{}) + forward := newHostedObservedWSConn(p.forward, p.forwardObservation) + bufferedForward := &hostedBufferedConn{Conn: forward, reader: bufio.NewReader(forward)} + forwardReadable := make(chan struct{}) + var forwardReady atomic.Bool + // Observe a lost forward channel while the supervisor is starting. The + // gateway must wait for readiness before sending HTTP/2; premature bytes + // close the lifetime. Hand off valid bytes only after Peek finishes. + go func() { + if _, err := bufferedForward.reader.Peek(1); err != nil || !forwardReady.Load() { + p.close() + } + close(forwardReadable) + }() + _ = p.reverse.SetReadDeadline(time.Time{}) + _ = p.reverse.SetWriteDeadline(time.Time{}) + reverse := newHostedObservedWSConn(p.reverse, p.reverseObservation) + client, err := newHostedHTTP2ClientConn(reverse) + if err != nil { + return + } + defer client.Close() //nolint:errcheck + go func() { + select { + case <-reverse.Done(): + p.close() + case <-p.done: + } + }() + for _, relay := range []struct { + address, target string + allow func(*http.Request) bool + }{ + {hostedBrokerRelayAddr, "http://broker", hostedBrokerRoute}, + {hostedOrkaRelayAddr, "http://orka", hostedOrkaRoute}, + } { + handler, err := newHostedProxy(relay.target, client, relay.allow) + if err != nil { + p.localFailure() + return + } + listener, err := net.Listen("tcp", relay.address) + if err != nil { + p.localFailure() + return + } + server := newHostedHTTPServer(handler) + defer server.Close() //nolint:errcheck + go func() { _ = server.Serve(listener) }() + } + sessionDir, err := os.MkdirTemp("/tmp", "orka-hosted-sessions-") + if err != nil || os.Chmod(sessionDir, 0o711) != nil { + p.localFailure() + return + } + // /tmp is outside the Foundry Session Files surface. Each wrapper lifetime + // launches exactly one supervisor; neither process is silently restarted. + env := make(map[string]string, len(p.bootstrap.Environment)+12) + for key, value := range p.bootstrap.Environment { + env[key] = value + } + for key, value := range map[string]string{ + "PATH": "/usr/local/sbin:/usr/local/bin:/usr/sbin:/usr/bin:/sbin:/bin", "HOME": "/tmp", + "ORKA_ACP_LISTEN_ADDRESS": hostedSupervisorAddr, + "ORKA_ACP_RUNTIME_INSTANCE_ID": "foundry-hosted." + s.challenge.BootID, + "ORKA_ACP_SUPERVISOR_BOOT_ID": s.challenge.BootID, + "ORKA_ACP_PROVIDER_PROXY_BASE_URL": "http://" + hostedBrokerRelayAddr + "/v1", + "ORKA_ACP_MCP_BROKER_URL": "http://" + hostedOrkaRelayAddr, + "ORKA_ACP_ARTIFACT_API_URL": "http://" + hostedOrkaRelayAddr, + "ORKA_ACP_WORKSPACE_MAX_ARTIFACT_BYTES": "536870912", + "ORKA_ACP_SESSION_BASE_DIR": sessionDir, + "ORKA_ACP_CONTROLLER_TOKEN_BOOTSTRAP": p.bootstrap.ControllerToken, + "ORKA_ACP_CAPABILITY_SECRET_BOOTSTRAP": p.bootstrap.CapabilitySecret, + "ORKA_ACP_PROVIDER_TOKEN_BOOTSTRAP": p.bootstrap.ProviderToken, + } { + env[key] = value + } + processDone, err := s.runner(s.ctx, env) + if err != nil { + if s.ctx.Err() == nil { + p.localFailure() + } + return + } + processResult := make(chan error, 1) + go func() { + err, ok := <-processDone + // Classify startup completion before close cancels the lifetime. A + // zero exit before health readiness is still a failed startup unless + // local or peer cancellation was already requested. + if !ok || (err == nil && !forwardReady.Load() && s.ctx.Err() == nil) { + err = errHostedInvalid + } + processResult <- err + p.close() + }() + defer p.joinSupervisor(processResult) + if waitHostedSupervisor(s.ctx) != nil { + if s.ctx.Err() == nil { + p.localFailure() + } + return + } + forwardReady.Store(true) + _ = p.forward.SetWriteDeadline(time.Now().Add(10 * time.Second)) + if p.forward.WriteJSON(hostedReady{Protocol: hostedProtocol, PairID: p.id, BootID: s.challenge.BootID, Ready: true}) != nil { + p.close() + return + } + _ = p.forward.SetWriteDeadline(time.Time{}) + transport := newHostedLocalTransport() + defer transport.CloseIdleConnections() + handler, err := newHostedProxy("http://"+hostedSupervisorAddr, transport, hostedV2Route) + if err != nil { + p.localFailure() + return + } + admission := http.HandlerFunc(func(w http.ResponseWriter, r *http.Request) { + if s.ctx.Err() != nil { + http.Error(w, "hosted lifetime closed", http.StatusServiceUnavailable) + return + } + handler.ServeHTTP(w, r) + }) + go func() { + select { + case <-forwardReadable: + serveHostedHTTP2(transportContext, bufferedForward, admission) + case <-transportContext.Done(): + } + p.close() + }() + select { + case <-s.ctx.Done(): + case <-forward.Done(): + p.close() + case <-reverse.Done(): + p.close() + } +} + +func (p *hostedPair) localFailure() { + p.server.mu.Lock() + p.shutdownErr = errHostedInvalid + p.server.mu.Unlock() +} + +func (p *hostedPair) joinSupervisor(result <-chan error) { + p.server.cancel() + defer p.close() + timer := time.NewTimer(hostedSupervisorShutdownWait) + defer timer.Stop() + var err error + select { + case err = <-result: + case <-timer.C: + err = errHostedInvalid + } + p.server.mu.Lock() + if err != nil { + p.shutdownErr = errHostedInvalid + } + p.server.mu.Unlock() +} + +type hostedBufferedConn struct { + net.Conn + reader *bufio.Reader +} + +func (c *hostedBufferedConn) Read(data []byte) (int, error) { return c.reader.Read(data) } + +func waitHostedSupervisor(ctx context.Context) error { + ctx, cancel := context.WithTimeout(ctx, 30*time.Second) + defer cancel() + transport := newHostedLocalTransport() + defer transport.CloseIdleConnections() + client := &http.Client{Transport: transport, Timeout: 2 * time.Second, + CheckRedirect: func(*http.Request, []*http.Request) error { return errHostedInvalid }} + for { + request, _ := http.NewRequestWithContext(ctx, http.MethodGet, "http://"+hostedSupervisorAddr+"/v2/health", nil) + response, err := client.Do(request) + if err == nil { + data, readErr := io.ReadAll(io.LimitReader(response.Body, 4097)) + _ = response.Body.Close() + var health struct { + Protocol string `json:"protocol"` + Status string `json:"status"` + } + if readErr == nil && len(data) <= 4096 && response.StatusCode == http.StatusOK && json.Unmarshal(data, &health) == nil && health.Protocol == "orka.harness.v2" && health.Status == "ok" { + return nil + } + } + select { + case <-ctx.Done(): + return errHostedInvalid + case <-time.After(100 * time.Millisecond): + } + } +} diff --git a/hosted_server_test.go b/hosted_server_test.go new file mode 100644 index 0000000..6bdd96e --- /dev/null +++ b/hosted_server_test.go @@ -0,0 +1,663 @@ +package main + +import ( + "bytes" + "context" + "encoding/json" + "errors" + "maps" + "net" + "net/http" + "net/http/httptest" + "os" + "path/filepath" + "reflect" + "strings" + "sync" + "sync/atomic" + "testing" + "time" + + "github.com/google/uuid" + "github.com/gorilla/websocket" + "golang.org/x/net/http2" +) + +type hostedServerTestFixture struct { + protocol hostedProtocolFixture + server *hostedServer + http *httptest.Server + calls atomic.Int32 + captured chan map[string]string + stopped chan struct{} + runnerOK bool + channelLogs *hostedTestChannelLog +} + +func newHostedServerTestFixture(t *testing.T, runnerOK bool) *hostedServerTestFixture { + t.Helper() + f := &hostedServerTestFixture{protocol: newHostedProtocolFixture(t), captured: make(chan map[string]string, 8), + stopped: make(chan struct{}, 8), runnerOK: runnerOK} + agent, err := json.Marshal(acpAgentConfiguration{Model: "test-model", ToolSchemaMode: toolSchemaModeRequest, + HostedTarget: f.protocol.config.Target}) + if err != nil { + t.Fatal("could not prepare hosted agent fixture") + } + f.protocol.config.AgentConfigurationDigest = brokerSHA(agent) + f.protocol.bootstrap.Environment["ORKA_ACP_AGENT_CONFIGURATION_DIGEST"] = brokerSHA(agent) + f.protocol.hello.Challenge.ConfigurationDigest = brokerJSONDigest(f.protocol.config) + f.protocol.hello.BootstrapDigest = brokerJSONDigest(f.protocol.bootstrap) + if validateHostedBootstrap(f.protocol.config, f.protocol.bootstrap) != nil { + t.Fatal("invalid canonical bootstrap fixture") + } + if _, err := decodeACPAgentConfiguration(agent, f.protocol.config.AgentConfigurationDigest, "test-model"); err != nil { + t.Fatal("invalid canonical agent configuration fixture") + } + ctx, cancel := context.WithCancel(context.Background()) + f.server = &hostedServer{ + cfg: f.protocol.config, challenge: f.protocol.hello.Challenge, agentConfig: agent, + adapterDigest: f.protocol.bootstrap.Environment["ORKA_ACP_FOUNDRY_ADAPTER_DIGEST"], ctx: ctx, cancel: cancel, + } + f.channelLogs = newHostedTestChannelLog() + f.server.channelLog = f.channelLogs.logger + f.server.runner = func(ctx context.Context, env map[string]string) (<-chan error, error) { + f.calls.Add(1) + f.captured <- maps.Clone(env) + if !f.runnerOK { + return nil, errors.New("fixture runner deliberately unavailable") + } + done := make(chan error, 1) + go func() { + <-ctx.Done() + done <- ctx.Err() + close(done) + f.stopped <- struct{}{} + }() + return done, nil + } + var handlers sync.WaitGroup + f.http = httptest.NewServer(http.HandlerFunc(func(w http.ResponseWriter, r *http.Request) { + handlers.Add(1) + defer handlers.Done() + f.server.ServeHTTP(w, r) + })) + t.Cleanup(func() { + cancel() + f.server.mu.Lock() + pair := f.server.pair + f.server.mu.Unlock() + if pair != nil { + pair.close() + } + f.http.Close() + finished := make(chan struct{}) + go func() { handlers.Wait(); close(finished) }() + hostedTestDone(t, finished) + for len(f.captured) > 0 { + env := <-f.captured + hostedServerRemoveTestDirectory(t, env["ORKA_ACP_SESSION_BASE_DIR"]) + } + if f.calls.Load() > 0 { + // Pair closure can precede run's deferred listener closes. Wait only + // for the two fixture-owned relay listeners before the next test. + for _, address := range []string{hostedBrokerRelayAddr, hostedOrkaRelayAddr} { + hostedServerWait(t, func() bool { + listener, err := net.Listen("tcp", address) + if err != nil { + return false + } + _ = listener.Close() + return true + }) + } + } + }) + return f +} + +func hostedServerRemoveTestDirectory(t *testing.T, path string) { + t.Helper() + if filepath.Dir(path) != "/tmp" || !strings.HasPrefix(filepath.Base(path), "orka-hosted-sessions-") { + t.Error("runner did not receive a dedicated temporary session directory") + return + } + if err := os.Remove(path); err != nil && !os.IsNotExist(err) { + t.Error("could not remove empty test session directory") + } +} + +func hostedServerWait(t *testing.T, condition func() bool) { + t.Helper() + deadline := time.NewTimer(3 * time.Second) + defer deadline.Stop() + tick := time.NewTicker(5 * time.Millisecond) + defer tick.Stop() + for { + if condition() { + return + } + select { + case <-deadline.C: + t.Fatal("hosted fixture state did not settle") + case <-tick.C: + } + } +} + +func (f *hostedServerTestFixture) connect(t *testing.T) *websocket.Conn { + t.Helper() + ws, response, err := (&websocket.Dialer{HandshakeTimeout: 3 * time.Second}).Dial( + "ws"+strings.TrimPrefix(f.http.URL, "http")+"/invocations_ws", nil) + if response != nil && response.Body != nil { + _ = response.Body.Close() + } + if err != nil { + t.Fatal("localhost hosted WebSocket upgrade failed") + } + t.Cleanup(func() { _ = ws.Close() }) + _ = ws.SetReadDeadline(time.Now().Add(3 * time.Second)) + var challenge hostedChallenge + if readHostedWSJSON(ws, &challenge) != nil || challenge != f.server.challenge { + t.Fatal("hosted server did not send its exact challenge") + } + return ws +} + +func (f *hostedServerTestFixture) hello(t *testing.T, role, pairID, digest string) hostedHello { + t.Helper() + return f.protocol.sign(t, hostedHello{Challenge: f.server.challenge, PairID: pairID, Role: role, + BootstrapDigest: digest, ExpiresAt: time.Now().Add(time.Minute).Unix()}) +} + +func (f *hostedServerTestFixture) sendBootstrap(t *testing.T, ws *websocket.Conn) { + t.Helper() + data, err := json.Marshal(f.protocol.bootstrap) + if err != nil || brokerSHA(data) != f.protocol.hello.BootstrapDigest || ws.WriteMessage(websocket.TextMessage, data) != nil { + t.Fatal("could not send exact signed bootstrap bytes") + } +} + +func (f *hostedServerTestFixture) claim(t *testing.T, role, pairID, digest string) *websocket.Conn { + t.Helper() + ws := f.connect(t) + if ws.WriteJSON(f.hello(t, role, pairID, digest)) != nil { + t.Fatal("could not send signed fixture hello") + } + var ack hostedAccepted + if readHostedWSJSON(ws, &ack) != nil || ack != (hostedAccepted{Protocol: hostedProtocol, + PairID: pairID, Role: role, BootID: f.server.challenge.BootID, BootstrapDigest: digest}) { + t.Fatal("authenticated role was not acknowledged exactly") + } + return ws +} + +func hostedServerRejected(t *testing.T, ws *websocket.Conn) { + t.Helper() + _ = ws.SetReadDeadline(time.Now().Add(3 * time.Second)) + _, _, err := ws.ReadMessage() + var networkError net.Error + if err == nil || errors.As(err, &networkError) && networkError.Timeout() { + t.Fatal("invalid handshake was acknowledged or left pending") + } +} + +func (f *hostedServerTestFixture) state(condition func(*hostedPair) bool) bool { + f.server.mu.Lock() + defer f.server.mu.Unlock() + return condition(f.server.pair) +} + +func (f *hostedServerTestFixture) receiveRunner(t *testing.T) map[string]string { + t.Helper() + select { + case env := <-f.captured: + t.Cleanup(func() { hostedServerRemoveTestDirectory(t, env["ORKA_ACP_SESSION_BASE_DIR"]) }) + return env + case <-time.After(3 * time.Second): + t.Fatal("authenticated pair did not reach the fixture runner") + return nil + } +} + +func hostedServerReverse(t *testing.T, ws *websocket.Conn) *hostedWSConn { + t.Helper() + conn := newHostedWSConn(ws) + ctx, cancel := context.WithCancel(context.Background()) + done := make(chan struct{}) + go func() { + defer close(done) + serveHostedHTTP2(ctx, conn, http.HandlerFunc(func(w http.ResponseWriter, _ *http.Request) { + w.WriteHeader(http.StatusNoContent) + })) + }() + t.Cleanup(func() { + cancel() + _ = conn.Close() + hostedTestDone(t, done) + }) + return conn +} + +func TestHostedServerUnauthenticatedInputCannotReservePair(t *testing.T) { + for _, name := range []string{"missing signature", "invalid signature", "different boot", "expired", "binary hello", "bootstrap before hello", "malformed JSON"} { + t.Run(name, func(t *testing.T) { + f := newHostedServerTestFixture(t, false) + ws := f.connect(t) + hello := f.hello(t, "forward", uuid.NewString(), brokerJSONDigest(f.protocol.bootstrap)) + kind := websocket.TextMessage + switch name { + case "missing signature": + hello.Signature = "" + case "invalid signature": + hello.Signature = strings.Repeat("a", len(hello.Signature)) + case "different boot": + hello.Challenge.BootID = uuid.NewString() + hello = f.protocol.sign(t, hello) + case "expired": + hello.ExpiresAt = time.Now().Add(-time.Second).Unix() + hello = f.protocol.sign(t, hello) + case "binary hello": + kind = websocket.BinaryMessage + } + data, _ := json.Marshal(hello) + if name == "bootstrap before hello" { + data, _ = json.Marshal(f.protocol.bootstrap) + } else if name == "malformed JSON" { + data = []byte("{") + } + if ws.WriteMessage(kind, data) != nil { + t.Fatal("could not send invalid admission fixture") + } + hostedServerRejected(t, ws) + if !f.state(func(p *hostedPair) bool { return p == nil }) || f.calls.Load() != 0 || f.server.ctx.Err() != nil { + t.Fatal("unauthenticated input reserved, launched or consumed the hosted lifetime") + } + }) + } +} + +func TestHostedServerConflictingClaimsCannotReplaceAuthenticatedRole(t *testing.T) { + for _, name := range []string{"forward replay", "different pair", "different bootstrap"} { + t.Run(name, func(t *testing.T) { + f := newHostedServerTestFixture(t, false) + pairID, digest := uuid.NewString(), brokerJSONDigest(f.protocol.bootstrap) + _ = f.claim(t, "forward", pairID, digest) + hostedServerWait(t, func() bool { return f.state(func(p *hostedPair) bool { return p != nil && p.forwardAcked }) }) + f.server.mu.Lock() + original := f.server.pair.forward + f.server.mu.Unlock() + role, otherPair, otherDigest := "reverse", pairID, digest + switch name { + case "forward replay": + role = "forward" + case "different pair": + otherPair = uuid.NewString() + case "different bootstrap": + otherDigest = brokerSHA([]byte("different bootstrap")) + } + ws := f.connect(t) + if ws.WriteJSON(f.hello(t, role, otherPair, otherDigest)) != nil { + t.Fatal("could not send conflicting fixture claim") + } + hostedServerRejected(t, ws) + if !f.state(func(p *hostedPair) bool { + return p.id == pairID && p.bootstrapDigest == digest && p.forward == original && p.reverse == nil && p.bootstrap == nil && !p.running + }) || f.calls.Load() != 0 || f.server.ctx.Err() != nil { + t.Fatal("conflicting claim replaced or disturbed admitted ownership") + } + }) + } +} + +func TestHostedServerConcurrentRoleReplayHasSingleWinner(t *testing.T) { + f := newHostedServerTestFixture(t, false) + pairID, digest := uuid.NewString(), brokerJSONDigest(f.protocol.bootstrap) + hello := f.hello(t, "forward", pairID, digest) + const count = 8 + results := make(chan *websocket.Conn, count) + for range count { + go func() { + ws, response, err := (&websocket.Dialer{HandshakeTimeout: 3 * time.Second}).Dial( + "ws"+strings.TrimPrefix(f.http.URL, "http")+"/invocations_ws", nil) + if response != nil && response.Body != nil { + _ = response.Body.Close() + } + if err != nil { + results <- nil + return + } + _ = ws.SetReadDeadline(time.Now().Add(3 * time.Second)) + var challenge hostedChallenge + var ack hostedAccepted + if readHostedWSJSON(ws, &challenge) == nil && challenge == hello.Challenge && ws.WriteJSON(hello) == nil && + readHostedWSJSON(ws, &ack) == nil && ack.PairID == pairID && ack.Role == "forward" { + results <- ws + return + } + _ = ws.Close() + results <- nil + }() + } + winners := 0 + for range count { + select { + case ws := <-results: + if ws != nil { + winners++ + t.Cleanup(func() { _ = ws.Close() }) + } + case <-time.After(5 * time.Second): + t.Fatal("concurrent claims did not finish") + } + } + if winners != 1 || f.calls.Load() != 0 || !f.state(func(p *hostedPair) bool { + return p != nil && p.id == pairID && p.forward != nil && p.reverse == nil && p.bootstrap == nil && !p.running + }) { + t.Fatal("concurrent replay admitted more than one role or launched early") + } +} + +func TestHostedServerRejectsMalformedOrUntrustedBootstrapBeforeLaunch(t *testing.T) { + for _, name := range []string{"wrong digest", "binary bootstrap", "unknown member", "duplicate member", "missing controller token", "model mismatch", "adapter mismatch", "inherited identity", "supervisor override", "invalid agent config"} { + t.Run(name, func(t *testing.T) { + f := newHostedServerTestFixture(t, false) + bootstrap := f.protocol.bootstrap + bootstrap.Environment = maps.Clone(bootstrap.Environment) + switch name { + case "missing controller token": + bootstrap.ControllerToken = "" + case "model mismatch": + bootstrap.Environment["ORKA_ACP_MODEL"] = "other-model" + case "adapter mismatch": + bootstrap.Environment["ORKA_ACP_FOUNDRY_ADAPTER_DIGEST"] = brokerSHA([]byte("other adapter")) + case "inherited identity": + bootstrap.Environment["IDENTITY_HEADER"] = "fixture-untrusted-identity" + case "supervisor override": + bootstrap.Environment["ORKA_ACP_SUPERVISOR_BOOT_ID"] = uuid.NewString() + case "invalid agent config": + f.server.agentConfig = append(bytes.Clone(f.server.agentConfig), '\n') + } + data, _ := json.Marshal(bootstrap) + if name == "unknown member" { + data = append([]byte(`{"injected":true,`), data[1:]...) + } else if name == "duplicate member" { + data = append([]byte(`{"environment":{},`), data[1:]...) + } + digest := brokerSHA(data) + if name == "wrong digest" { + digest = brokerSHA([]byte("other bytes")) + } + pairID := uuid.NewString() + forward := f.claim(t, "forward", pairID, digest) + _ = f.claim(t, "reverse", pairID, digest) + kind := websocket.TextMessage + if name == "binary bootstrap" { + kind = websocket.BinaryMessage + } + if forward.WriteMessage(kind, data) != nil { + t.Fatal("could not send invalid bootstrap fixture") + } + hostedTestDone(t, f.server.ctx.Done()) + if f.calls.Load() != 0 || !f.state(func(p *hostedPair) bool { return p != nil && p.bootstrap == nil && !p.running }) { + t.Fatal("untrusted bootstrap reached the supervisor runner") + } + }) + } +} + +func TestHostedServerRequiresBothRolesAndExactConstructedEnvironment(t *testing.T) { + for _, order := range []string{"forward first", "reverse first"} { + t.Run(order, func(t *testing.T) { + for _, name := range []string{"IDENTITY_HEADER", "IDENTITY_ENDPOINT", "AZURE_CLIENT_ID", "AZURE_TENANT_ID", "AZURE_CLIENT_SECRET", "AZURE_FEDERATED_TOKEN_FILE", "MSI_SECRET"} { + t.Setenv(name, "fixture-parent-only") + } + f := newHostedServerTestFixture(t, false) + pairID, digest := uuid.NewString(), brokerJSONDigest(f.protocol.bootstrap) + var forward, reverse *websocket.Conn + if order == "forward first" { + forward = f.claim(t, "forward", pairID, digest) + f.sendBootstrap(t, forward) + hostedServerWait(t, func() bool { return f.state(func(p *hostedPair) bool { return p.bootstrap != nil }) }) + } else { + reverse = f.claim(t, "reverse", pairID, digest) + hostedServerWait(t, func() bool { return f.state(func(p *hostedPair) bool { return p.reverseAcked }) }) + } + if f.calls.Load() != 0 || !f.state(func(p *hostedPair) bool { return p != nil && !p.running }) { + t.Fatal("a single authenticated role launched the supervisor") + } + if order == "forward first" { + reverse = f.claim(t, "reverse", pairID, digest) + } else { + forward = f.claim(t, "forward", pairID, digest) + hostedServerWait(t, func() bool { return f.state(func(p *hostedPair) bool { return p.forwardAcked && p.reverseAcked }) }) + if f.calls.Load() != 0 || !f.state(func(p *hostedPair) bool { return p.bootstrap == nil && !p.running }) { + t.Fatal("authenticated roles launched without the signed bootstrap") + } + } + _ = hostedServerReverse(t, reverse) + if order == "reverse first" { + f.sendBootstrap(t, forward) + } + env := f.receiveRunner(t) + expected := maps.Clone(f.protocol.bootstrap.Environment) + for key, value := range map[string]string{ + "PATH": "/usr/local/sbin:/usr/local/bin:/usr/sbin:/usr/bin:/sbin:/bin", "HOME": "/tmp", + "ORKA_ACP_LISTEN_ADDRESS": hostedSupervisorAddr, + "ORKA_ACP_RUNTIME_INSTANCE_ID": "foundry-hosted." + f.server.challenge.BootID, + "ORKA_ACP_SUPERVISOR_BOOT_ID": f.server.challenge.BootID, + "ORKA_ACP_PROVIDER_PROXY_BASE_URL": "http://" + hostedBrokerRelayAddr + "/v1", + "ORKA_ACP_MCP_BROKER_URL": "http://" + hostedOrkaRelayAddr, + "ORKA_ACP_ARTIFACT_API_URL": "http://" + hostedOrkaRelayAddr, + "ORKA_ACP_WORKSPACE_MAX_ARTIFACT_BYTES": "536870912", + "ORKA_ACP_SESSION_BASE_DIR": env["ORKA_ACP_SESSION_BASE_DIR"], + "ORKA_ACP_CONTROLLER_TOKEN_BOOTSTRAP": f.protocol.bootstrap.ControllerToken, + "ORKA_ACP_CAPABILITY_SECRET_BOOTSTRAP": f.protocol.bootstrap.CapabilitySecret, + "ORKA_ACP_PROVIDER_TOKEN_BOOTSTRAP": f.protocol.bootstrap.ProviderToken, + } { + expected[key] = value + } + if !reflect.DeepEqual(env, expected) { + t.Fatal("supervisor environment inherited authority or changed its pinned configuration") + } + info, err := os.Lstat(env["ORKA_ACP_SESSION_BASE_DIR"]) + if err != nil || !info.IsDir() || info.Mode().Perm() != 0o711 { + t.Fatal("supervisor session directory lacks its explicit isolation permissions") + } + hostedTestDone(t, f.server.ctx.Done()) + if f.calls.Load() != 1 { + t.Fatal("fixture runner failure retried the supervisor") + } + }) + } +} + +func hostedServerTestHealth(t *testing.T, status int) <-chan struct{} { + t.Helper() + called := make(chan struct{}, 1) + listener, err := net.Listen("tcp", hostedSupervisorAddr) + if err != nil { + t.Fatal("fixed supervisor fixture port is unavailable") + } + server := httptest.NewUnstartedServer(http.HandlerFunc(func(w http.ResponseWriter, r *http.Request) { + if r.Method != http.MethodGet || r.URL.Path != "/v2/health" { + http.NotFound(w, r) + return + } + select { + case called <- struct{}{}: + default: + } + w.Header().Set("Content-Type", "application/json") + w.WriteHeader(status) + _, _ = w.Write([]byte(`{"protocol":"orka.harness.v2","status":"ok"}`)) + })) + _ = server.Listener.Close() + server.Listener = listener + server.Start() + t.Cleanup(server.Close) + return called +} + +func hostedServerTestRunningPair(t *testing.T, f *hostedServerTestFixture) (*hostedWSConn, *hostedWSConn, *http2.ClientConn) { + t.Helper() + pairID, digest := uuid.NewString(), brokerJSONDigest(f.protocol.bootstrap) + forward := f.claim(t, "forward", pairID, digest) + reverse := hostedServerReverse(t, f.claim(t, "reverse", pairID, digest)) + f.sendBootstrap(t, forward) + _ = f.receiveRunner(t) + var ready hostedReady + if readHostedWSJSON(forward, &ready) != nil || ready != (hostedReady{Protocol: hostedProtocol, + PairID: pairID, BootID: f.server.challenge.BootID, Ready: true}) { + t.Fatal("running fixture did not report exact readiness") + } + forwardConn := newHostedWSConn(forward) + client, err := newHostedHTTP2ClientConn(forwardConn) + if err != nil { + t.Fatal("could not establish forward HTTP/2 fixture") + } + t.Cleanup(func() { _ = client.Close(); _ = forwardConn.Close() }) + return forwardConn, reverse, client +} + +func TestHostedServerRunningLifetimeCannotLaunchSecondSupervisor(t *testing.T) { + hostedServerTestHealth(t, http.StatusOK) + f := newHostedServerTestFixture(t, true) + _, _, client := hostedServerTestRunningPair(t, f) + ctx, cancel := context.WithTimeout(context.Background(), 3*time.Second) + defer cancel() + request, _ := http.NewRequestWithContext(ctx, http.MethodGet, "http://supervisor/v2/health", nil) + response, err := client.RoundTrip(request) + if err != nil { + t.Fatal("authenticated pair did not carry a forward HTTP/2 request") + } + _ = response.Body.Close() + if response.StatusCode != http.StatusOK { + t.Fatal("forward proxy changed supervisor health status") + } + for _, role := range []string{"forward", "reverse"} { + ws := f.connect(t) + if ws.WriteJSON(f.hello(t, role, uuid.NewString(), brokerJSONDigest(f.protocol.bootstrap))) != nil { + t.Fatal("could not send second-launch fixture") + } + hostedServerRejected(t, ws) + } + if f.calls.Load() != 1 || f.server.ctx.Err() != nil { + t.Fatal("second launch was admitted or disrupted the authenticated lifetime") + } +} + +func TestHostedServerEitherChannelLossClosesLifetime(t *testing.T) { + for _, role := range []string{"forward", "reverse"} { + t.Run(role, func(t *testing.T) { + hostedServerTestHealth(t, http.StatusOK) + f := newHostedServerTestFixture(t, true) + forward, reverse, _ := hostedServerTestRunningPair(t, f) + if role == "forward" { + _ = forward.Close() + } else { + _ = reverse.Close() + } + hostedTestDone(t, f.server.ctx.Done()) + hostedTestDone(t, forward.Done()) + hostedTestDone(t, reverse.Done()) + select { + case <-f.stopped: + case <-time.After(3 * time.Second): + t.Fatal("channel loss did not cancel the supervisor runner") + } + response := httptest.NewRecorder() + f.server.ServeHTTP(response, httptest.NewRequest(http.MethodGet, "/readiness", nil)) + if response.Code != http.StatusServiceUnavailable || f.calls.Load() != 1 { + t.Fatal("closed lifetime remained ready or replayed its supervisor") + } + ws, reply, err := (&websocket.Dialer{HandshakeTimeout: 3 * time.Second}).Dial( + "ws"+strings.TrimPrefix(f.http.URL, "http")+"/invocations_ws", nil) + if ws != nil { + _ = ws.Close() + } + if reply != nil && reply.Body != nil { + _ = reply.Body.Close() + } + if err == nil || reply == nil || reply.StatusCode != http.StatusServiceUnavailable { + t.Fatal("closed lifetime accepted a replayed channel") + } + }) + } +} + +func TestHostedServerIdleGoingAwayClosesLifetime(t *testing.T) { + for _, role := range []string{"forward", "reverse"} { + t.Run(role, func(t *testing.T) { + hostedServerTestHealth(t, http.StatusOK) + f := newHostedServerTestFixture(t, true) + earliest := time.Now() + forward, reverse, _ := hostedServerTestRunningPair(t, f) + latest := time.Now() + channel := forward + if role == "reverse" { + channel = reverse + } + hostedTestGoingAway(t, channel.ws) + hostedTestDone(t, f.server.ctx.Done()) + hostedTestDone(t, forward.Done()) + hostedTestDone(t, reverse.Done()) + f.server.mu.Lock() + stopped := f.server.pair.stopped + f.server.mu.Unlock() + hostedTestDone(t, stopped) + response := httptest.NewRecorder() + f.server.ServeHTTP(response, httptest.NewRequest(http.MethodGet, "/readiness", nil)) + if response.Code != http.StatusServiceUnavailable || f.calls.Load() != 1 { + t.Fatal("idle platform closure left a ready runtime or relaunched its supervisor") + } + replay := httptest.NewRecorder() + f.server.ServeHTTP(replay, httptest.NewRequest(http.MethodGet, "/invocations_ws", nil)) + if replay.Code != http.StatusServiceUnavailable || f.calls.Load() != 1 { + t.Fatal("platform-closed lifetime admitted a replacement channel") + } + hostedTestClosedChannelLogs(t, f.channelLogs, role, earliest, latest) + }) + } +} + +func TestHostedServerEitherChannelLossWhileStartingCancelsRunner(t *testing.T) { + for _, role := range []string{"forward", "reverse", "forward after early preface"} { + t.Run(role, func(t *testing.T) { + healthCalled := hostedServerTestHealth(t, http.StatusServiceUnavailable) + f := newHostedServerTestFixture(t, true) + pairID, digest := uuid.NewString(), brokerJSONDigest(f.protocol.bootstrap) + forward := f.claim(t, "forward", pairID, digest) + reverse := hostedServerReverse(t, f.claim(t, "reverse", pairID, digest)) + f.sendBootstrap(t, forward) + _ = f.receiveRunner(t) + select { + case <-healthCalled: + case <-time.After(3 * time.Second): + t.Fatal("fixture did not reach supervisor readiness polling") + } + if role == "reverse" { + _ = reverse.Close() + } else { + if role == "forward after early preface" { + if forward.WriteMessage(websocket.BinaryMessage, []byte(http2.ClientPreface)) != nil { + t.Fatal("could not send the early HTTP/2 preface") + } + } + _ = forward.Close() + } + select { + case <-f.server.ctx.Done(): + case <-time.After(time.Second): + t.Fatal("channel loss left the starting supervisor alive") + } + select { + case <-f.stopped: + case <-time.After(3 * time.Second): + t.Fatal("lost startup channel did not cancel the runner") + } + if f.calls.Load() != 1 { + t.Fatal("lost startup channel replayed the supervisor") + } + }) + } +} diff --git a/hosted_startup_config_test.go b/hosted_startup_config_test.go new file mode 100644 index 0000000..e151ed9 --- /dev/null +++ b/hosted_startup_config_test.go @@ -0,0 +1,89 @@ +package main + +import ( + "bytes" + "crypto/ed25519" + "encoding/base64" + "encoding/json" + "maps" + "os" + "path/filepath" + "strings" + "testing" +) + +func TestHostedGatewayRejectsUnusableSigningKeyDuringSettingsLoad(t *testing.T) { + for _, scenario := range []string{"valid", "zero-seed", "mismatched-seed"} { + t.Run(scenario, func(t *testing.T) { + f := newHostedProtocolFixture(t) + key := bytes.Clone(f.key) + defer clear(key) + switch scenario { + case "zero-seed": + key = ed25519.NewKeyFromSeed(make([]byte, ed25519.SeedSize)) + defer clear(key) + f.config.SigningPublicKey = base64.RawURLEncoding.EncodeToString(key[ed25519.SeedSize:]) + case "mismatched-seed": + key[0] ^= 1 + } + cfg := hostedGatewayConfig{ + Protocol: hostedProtocol, Image: f.config, + ContainerImage: "example.invalid/hosted@" + brokerSHA([]byte("fixture image")), + SessionID: f.hello.Challenge.SessionID, RuntimeProfileDigest: brokerSHA([]byte("fixture profile")), + RuntimeEnvironment: maps.Clone(f.bootstrap.Environment), + OrkaBaseURL: "http://orka.test:8080", BrokerBaseURL: "http://127.0.0.1:8091", + } + if validateHostedGatewayConfig(cfg) != nil { + t.Fatal("signing-key fixture has invalid public configuration") + } + dir := t.TempDir() + write := func(name string, data []byte) string { + t.Helper() + path := filepath.Join(dir, name) + if os.WriteFile(path, data, 0o600) != nil { + t.Fatal("could not write local settings fixture") + } + return path + } + raw, err := json.Marshal(cfg) + if err != nil { + t.Fatal("could not marshal public settings fixture") + } + configPath := write("gateway.json", raw) + env := map[string]string{ + "ORKA_FOUNDRY_GATEWAY_STATE_DIR": filepath.Join(dir, "ledger"), + "ORKA_FOUNDRY_GATEWAY_SIGNING_KEY_FILE": write("signing-key", []byte(base64.RawURLEncoding.EncodeToString(key))), + "ORKA_FOUNDRY_GATEWAY_CONTROLLER_TOKEN_FILE": write("controller-token", []byte(f.bootstrap.ControllerToken)), + "ORKA_FOUNDRY_GATEWAY_CAPABILITY_SECRET_FILE": write("capability-secret", []byte(f.bootstrap.CapabilitySecret)), + "ORKA_FOUNDRY_GATEWAY_PROVIDER_TOKEN_FILE": write("provider-token", []byte(f.bootstrap.ProviderToken)), + } + bootstrapReads := 0 + settings, err := loadHostedGatewaySettings(configPath, func(name string) string { + if strings.HasSuffix(name, "_TOKEN_FILE") || strings.HasSuffix(name, "_SECRET_FILE") { + bootstrapReads++ + } + return env[name] + }) + defer clear(settings.signingKey) + if scenario != "valid" { + if err == nil { + t.Fatal("unusable signing key passed the pre-network settings boundary") + } + if bootstrapReads != 0 || hostedNonzeroBytes(settings.signingKey) { + t.Fatal("invalid signing material was retained or bootstrap credentials were read") + } + return + } + if err != nil || bootstrapReads != 3 { + t.Fatal("valid signing-key configuration was rejected") + } + signed, err := signHostedHello(f.hello, settings.signingKey) + if err != nil || verifyHostedHello(signed, f.hello.Challenge, cfg.Image.SigningPublicKey, f.now) != nil { + t.Fatal("loaded signing key could not authenticate the configured hosted image") + } + if _, err := os.Stat(env["ORKA_FOUNDRY_GATEWAY_STATE_DIR"]); !os.IsNotExist(err) { + t.Fatal("settings validation created a gateway ownership ledger") + } + }) + } +} diff --git a/hosted_startup_exit_test.go b/hosted_startup_exit_test.go new file mode 100644 index 0000000..60d8370 --- /dev/null +++ b/hosted_startup_exit_test.go @@ -0,0 +1,98 @@ +package main + +import ( + "context" + "errors" + "maps" + "net/http" + "net/http/httptest" + "testing" + "time" + + "github.com/google/uuid" +) + +func TestHostedStartupProcessCompletionBeforeReadiness(t *testing.T) { + for _, cause := range []string{"unexpected-nil", "unexpected-error", "missing-result", "local-cancel", "forward-close", "reverse-close"} { + t.Run(cause, func(t *testing.T) { + healthCalled := hostedServerTestHealth(t, http.StatusServiceUnavailable) + f := newHostedServerTestFixture(t, true) + processDone := make(chan error, 1) + f.server.runner = func(_ context.Context, environment map[string]string) (<-chan error, error) { + f.calls.Add(1) + f.captured <- maps.Clone(environment) + return processDone, nil + } + pairID, digest := uuid.NewString(), brokerJSONDigest(f.protocol.bootstrap) + forward := f.claim(t, "forward", pairID, digest) + reverse := hostedServerReverse(t, f.claim(t, "reverse", pairID, digest)) + f.sendBootstrap(t, forward) + _ = f.receiveRunner(t) + hostedTestDone(t, healthCalled) + if f.server.ctx.Err() != nil { + t.Fatal("fixture lifetime closed before the supervisor exit") + } + returned := make(chan error, 1) + go func() { returned <- serveHostedHTTP(f.server.ctx, "127.0.0.1:0", f.server) }() + alreadyCancelled := false + switch cause { + case "local-cancel": + f.server.cancel() + alreadyCancelled = true + case "forward-close": + hostedTestGoingAway(t, forward) + alreadyCancelled = true + case "reverse-close": + hostedTestGoingAway(t, reverse.ws) + alreadyCancelled = true + } + if alreadyCancelled { + hostedTestDone(t, f.server.ctx.Done()) + select { + case <-returned: + t.Fatal("entrypoint returned before the cancelled supervisor was joined") + default: + } + } + switch cause { + case "missing-result": + close(processDone) + case "unexpected-error": + processDone <- errHostedInvalid + close(processDone) + default: + processDone <- nil + close(processDone) + } + select { + case err := <-returned: + if alreadyCancelled { + if err != nil { + t.Error("already-cancelled startup was reported as local process failure") + } + } else if !errors.Is(err, errHostedInvalid) { + t.Error("supervisor exited before health became ready but the entrypoint reported success") + } + case <-time.After(3 * time.Second): + t.Fatal("entrypoint did not join the completed supervisor") + } + f.server.mu.Lock() + pair := f.server.pair + f.server.mu.Unlock() + hostedTestDone(t, pair.stopped) + hostedTestDone(t, pair.done) + hostedTestDone(t, reverse.Done()) + hostedServerRejected(t, forward) + for _, path := range []string{"/readiness", "/invocations_ws"} { + reply := httptest.NewRecorder() + f.server.ServeHTTP(reply, httptest.NewRequest(http.MethodGet, path, nil)) + if reply.Code != http.StatusServiceUnavailable { + t.Error("completed startup retained readiness or admitted another channel") + } + } + if f.calls.Load() != 1 { + t.Error("completed startup relaunched its supervisor") + } + }) + } +} diff --git a/hosted_startup_review_test.go b/hosted_startup_review_test.go new file mode 100644 index 0000000..6e4a5f1 --- /dev/null +++ b/hosted_startup_review_test.go @@ -0,0 +1,304 @@ +package main + +import ( + "bytes" + "context" + "errors" + "maps" + "net" + "net/http" + "os" + "path/filepath" + "sync" + "sync/atomic" + "testing" + "time" + + "github.com/google/uuid" +) + +func hostedStartupGatewaySettings(t *testing.T) hostedGatewaySettings { + t.Helper() + f := newHostedProtocolFixture(t) + settings := hostedGatewaySettings{ + config: hostedGatewayConfig{ + Protocol: hostedProtocol, Image: f.config, + ContainerImage: "example.invalid/hosted@" + brokerSHA([]byte("startup fixture image")), + SessionID: f.hello.Challenge.SessionID, RuntimeProfileDigest: brokerSHA([]byte("startup fixture profile")), + RuntimeEnvironment: maps.Clone(f.bootstrap.Environment), + OrkaBaseURL: "http://orka.test:8080", BrokerBaseURL: "http://127.0.0.1:8091", + }, + bootstrap: f.bootstrap, signingKey: bytes.Clone(f.key), stateDir: filepath.Join(t.TempDir(), "gateway"), + } + t.Cleanup(func() { clear(settings.signingKey) }) + if validateHostedGatewayConfig(settings.config) != nil { + t.Fatal("invalid gateway startup fixture") + } + return settings +} + +func TestHostedStartupOccupiedAddressDoesNotInitializeGateway(t *testing.T) { + reserved, err := net.Listen("tcp", "127.0.0.1:0") + if err != nil { + t.Fatal("could not reserve fixture address") + } + defer reserved.Close() //nolint:errcheck + settings := hostedStartupGatewaySettings(t) + settings.address = reserved.Addr().String() + var tokens atomic.Int32 + provider := hostedGatewayTestTokenProvider(func(context.Context) (string, error) { + tokens.Add(1) + // No token is returned, so this fixture cannot make a remote call. + return "", errHostedInvalid + }) + if !errors.Is(serveHostedGateway(t.Context(), settings, provider), errHostedInvalid) { + t.Fatal("occupied listener did not reject startup") + } + if tokens.Load() != 0 { + t.Error("occupied listener reached gateway token acquisition") + } + if _, err := os.Stat(settings.stateDir); !os.IsNotExist(err) { + t.Error("occupied listener created a gateway ownership ledger") + } + if hostedNonzeroBytes(settings.signingKey) { + t.Error("failed listener retained signing material") + } +} + +func TestHostedStartupListenerHeldThroughInitializationAndReleasedOnFailure(t *testing.T) { + reserved, err := net.Listen("tcp", "127.0.0.1:0") + if err != nil { + t.Fatal("could not select fixture address") + } + settings := hostedStartupGatewaySettings(t) + settings.address = reserved.Addr().String() + _ = reserved.Close() + var tokens atomic.Int32 + provider := hostedGatewayTestTokenProvider(func(context.Context) (string, error) { + tokens.Add(1) + probe, err := net.Listen("tcp", settings.address) + if err == nil { + _ = probe.Close() + t.Error("gateway initialized before acquiring its listener") + } + return "", errHostedInvalid + }) + if !errors.Is(serveHostedGateway(t.Context(), settings, provider), errHostedInvalid) || tokens.Load() != 1 { + t.Fatal("fixture did not reach the intended initialization failure") + } + listener, err := net.Listen("tcp", settings.address) + if err != nil { + t.Fatal("failed initialization retained the gateway listener") + } + _ = listener.Close() + if hostedNonzeroBytes(settings.signingKey) { + t.Error("failed initialization retained signing material") + } +} + +func TestHostedStartupServesRetainedListener(t *testing.T) { + listener, err := net.Listen("tcp", "127.0.0.1:0") + if err != nil { + t.Fatal("could not bind fixture listener") + } + defer listener.Close() //nolint:errcheck + ctx, cancel := context.WithCancel(t.Context()) + defer cancel() + done := make(chan error, 1) + go func() { + done <- serveHostedHTTPListener(ctx, listener, http.HandlerFunc(func(w http.ResponseWriter, _ *http.Request) { + w.WriteHeader(http.StatusNoContent) + })) + }() + transport := newHostedLocalTransport() + defer transport.CloseIdleConnections() + client := &http.Client{Transport: transport, Timeout: 2 * time.Second} + response, err := client.Get("http://" + listener.Addr().String() + "/healthz") + if err != nil { + t.Fatal("retained listener did not serve HTTP") + } + _ = response.Body.Close() + if response.StatusCode != http.StatusNoContent { + t.Error("retained listener changed the handler response") + } + cancel() + select { + case err := <-done: + if err != nil { + t.Error("retained listener did not shut down cleanly") + } + case <-time.After(3 * time.Second): + t.Fatal("retained listener did not close after cancellation") + } + if _, err := listener.Accept(); !errors.Is(err, net.ErrClosed) { + t.Fatal("cancelled server retained its listener") + } +} + +func TestHostedStartupSessionPreflightFailureDoesNotReserveCreate(t *testing.T) { + for _, failure := range []string{"token-error", "malformed-token", "principal-drift", "cancelled-token", "cancelled-after-token"} { + t.Run(failure, func(t *testing.T) { + f := newHostedGatewayTestFixture(t, hostedGatewayTestOptions{}) + var submissions atomic.Int32 + transport := f.gateway.httpClient.Transport + f.gateway.httpClient.Transport = hostedGatewayTestRoundTripper(func(request *http.Request) (*http.Response, error) { + if request.Method == http.MethodPost { + submissions.Add(1) + } + return transport.RoundTrip(request) + }) + f.gateway.provider = hostedGatewayTestTokenProvider(func(ctx context.Context) (string, error) { + call := f.tokenCalls.Add(1) + if call == 4 { // Agent, version and exact-session GETs precede creation. + switch failure { + case "token-error": + return "", errHostedInvalid + case "malformed-token": + return "invalid-fixture-identity", nil + case "principal-drift": + claims := maps.Clone(f.claims) + claims["oid"] = uuid.NewString() + return hostedGatewayTestJWT(claims), nil + case "cancelled-token": + f.gateway.cancel() + return "", context.Canceled + case "cancelled-after-token": + f.gateway.cancel() + } + } + if ctx.Err() != nil && failure != "cancelled-after-token" { + return "", ctx.Err() + } + return f.token.Load().(string), nil + }) + if f.initialize() == nil || f.tokenCalls.Load() != 4 { + t.Fatal("fixture did not reject the create preflight") + } + if submissions.Load() != 0 || f.creates.Load() != 0 || f.dials.Load() != 0 || f.httpCalls.Load() != 3 { + t.Error("failed create preflight reached transport submission or channel setup") + } + f.assertNoBootstrap(t) + store, ledger, err := openHostedGatewayStore(f.settings.stateDir, f.settings.config) + if err != nil { + t.Fatal("could not reopen the preflight fixture ledger") + } + defer store.close() + if ledger.CreateAttempted || ledger.SessionCreated || ledger.ExposurePossible || ledger.PrincipalDigest == "" { + t.Fatal("definitely-unsent creation retained ambiguous ownership or lost its principal") + } + // A new startup may use the same pinned principal after a preflight + // failure. It still performs the exact-session GET before one POST. + restarted := &hostedGateway{cfg: f.settings.config, provider: f.gateway.provider, + httpClient: f.gateway.httpClient, store: store, ledger: ledger} + if restarted.ensureSession(t.Context()) != nil || submissions.Load() != 1 || f.creates.Load() != 1 || !restarted.ledger.SessionCreated { + t.Fatal("definitely-unsent startup could not make one later owned creation") + } + }) + } +} + +func TestHostedStartupCreateReservationPrecedesSubmission(t *testing.T) { + f := newHostedGatewayTestFixture(t, hostedGatewayTestOptions{}) + transport := f.gateway.httpClient.Transport + var submissions atomic.Int32 + f.gateway.httpClient.Transport = hostedGatewayTestRoundTripper(func(request *http.Request) (*http.Response, error) { + if request.Method == http.MethodPost { + submissions.Add(1) + data, err := readHostedFile(filepath.Join(f.settings.stateDir, "state.json"), hostedMaxHandshakeBytes) + var ledger hostedGatewayLedger + if err != nil || acpDecode(data, &ledger, true) != nil || !ledger.CreateAttempted || ledger.SessionCreated || ledger.ExposurePossible || ledger.PrincipalDigest == "" { + t.Error("session creation reached transport before durable ownership") + } + if request.GetBody != nil { + t.Error("session creation supplied a replayable request body") + } + } + return transport.RoundTrip(request) + }) + if f.gateway.ensureSession(t.Context()) != nil || submissions.Load() != 1 || f.creates.Load() != 1 { + t.Fatal("owned session creation did not complete exactly once") + } +} + +func TestHostedStartupSubmittedCreateErrorRemainsReserved(t *testing.T) { + f := newHostedGatewayTestFixture(t, hostedGatewayTestOptions{}) + transport := f.gateway.httpClient.Transport + var submissions atomic.Int32 + f.gateway.httpClient.Transport = hostedGatewayTestRoundTripper(func(request *http.Request) (*http.Response, error) { + if request.Method == http.MethodPost { + submissions.Add(1) + return nil, context.Canceled + } + return transport.RoundTrip(request) + }) + if f.initialize() == nil || submissions.Load() != 1 || f.creates.Load() != 0 { + t.Fatal("fixture did not fail at transport submission") + } + store, ledger, err := openHostedGatewayStore(f.settings.stateDir, f.settings.config) + if err != nil { + t.Fatal("could not reopen submitted-create ledger") + } + defer store.close() + if !ledger.CreateAttempted || ledger.SessionCreated || ledger.ExposurePossible { + t.Fatal("transport error erased an uncertain creation") + } + restarted := &hostedGateway{cfg: f.settings.config, provider: f.gateway.provider, + httpClient: f.gateway.httpClient, store: store, ledger: ledger} + tokens, requests := f.tokenCalls.Load(), f.httpCalls.Load() + if restarted.ensureSession(t.Context()) == nil || f.tokenCalls.Load() != tokens || f.httpCalls.Load() != requests || submissions.Load() != 1 { + t.Fatal("uncertain creation was replayed or adopted after restart") + } +} + +func TestHostedStartupRunnerFailureClassification(t *testing.T) { + for _, cause := range []string{"uncancelled-error", "local-cancel", "forward-close", "reverse-close"} { + t.Run(cause, func(t *testing.T) { + f := newHostedServerTestFixture(t, false) + entered, released := make(chan struct{}), make(chan struct{}) + release := sync.OnceFunc(func() { close(released) }) + t.Cleanup(release) + var cancelled atomic.Bool + f.server.runner = func(ctx context.Context, environment map[string]string) (<-chan error, error) { + f.calls.Add(1) + f.captured <- maps.Clone(environment) + close(entered) + <-released + cancelled.Store(ctx.Err() != nil) + // Match the production runner's rejection before command.Start. + return nil, errHostedInvalid + } + pairID, digest := uuid.NewString(), brokerJSONDigest(f.protocol.bootstrap) + forward := f.claim(t, "forward", pairID, digest) + reverse := hostedServerReverse(t, f.claim(t, "reverse", pairID, digest)) + f.sendBootstrap(t, forward) + hostedTestDone(t, entered) + switch cause { + case "local-cancel": + f.server.cancel() + case "forward-close": + hostedTestGoingAway(t, forward) + case "reverse-close": + hostedTestGoingAway(t, reverse.ws) + } + if cause != "uncancelled-error" { + hostedTestDone(t, f.server.ctx.Done()) + } + release() + f.server.mu.Lock() + pair := f.server.pair + f.server.mu.Unlock() + hostedTestDone(t, pair.stopped) + err := f.server.close() + if cause == "uncancelled-error" { + if !errors.Is(err, errHostedInvalid) || cancelled.Load() { + t.Error("actual runner startup failure was reported as clean cancellation") + } + } else if err != nil || !cancelled.Load() { + t.Error("cancellation before runner launch was reported as a local failure") + } + if f.calls.Load() != 1 { + t.Error("failed or cancelled runner launch was retried") + } + }) + } +} diff --git a/hosted_store.go b/hosted_store.go new file mode 100644 index 0000000..6d68e90 --- /dev/null +++ b/hosted_store.go @@ -0,0 +1,118 @@ +package main + +import ( + "encoding/json" + "io" + "os" + "path/filepath" + "syscall" + + "github.com/google/uuid" +) + +// The gateway records possible credential exposure BEFORE sending bootstrap. +// A replacement process cannot seed a fresh supervisor from a nonempty record. +// This deliberately blocks unknown ownership instead of replaying mutations. +type hostedGatewayLedger struct { + Version uint32 `json:"version"` + ConfigDigest string `json:"configDigest"` + SessionID string `json:"sessionID"` + PrincipalDigest string `json:"principalDigest,omitempty"` + CreateAttempted bool `json:"createAttempted"` + SessionCreated bool `json:"sessionCreated"` + ExposurePossible bool `json:"exposurePossible"` + Challenge hostedChallenge `json:"challenge"` + PairID string `json:"pairID,omitempty"` + BootstrapDigest string `json:"bootstrapDigest,omitempty"` + Ready bool `json:"ready"` + Closed bool `json:"closed"` +} + +type hostedGatewayStore struct { + dir string + lock *os.File +} + +func openHostedGatewayStore(dir string, cfg hostedGatewayConfig) (*hostedGatewayStore, hostedGatewayLedger, error) { + var empty hostedGatewayLedger + if !filepath.IsAbs(dir) || filepath.Clean(dir) == "/" { + return nil, empty, errHostedInvalid + } + lock, created, err := openStoreLock(dir, "gateway.lock", syncStoreDirectory) + if err != nil { + return nil, empty, errHostedInvalid + } + store := &hostedGatewayStore{dir: dir, lock: lock} + ledger := hostedGatewayLedger{Version: 1, ConfigDigest: brokerJSONDigest(cfg), SessionID: cfg.SessionID} + path := filepath.Join(dir, "state.json") + info, err := os.Lstat(path) + if os.IsNotExist(err) { + if created && store.save(ledger) == nil { + return store, ledger, nil + } + } else if err == nil && brokerPrivateFile(info) && info.Size() <= hostedMaxHandshakeBytes { + data, readErr := readHostedFile(path, hostedMaxHandshakeBytes) + var stored hostedGatewayLedger + if readErr == nil && acpDecode(data, &stored, true) == nil && hostedGatewayLedgerValid(stored, cfg) { + return store, stored, nil + } + } + store.close() + return nil, empty, errHostedInvalid +} + +func hostedGatewayLedgerValid(value hostedGatewayLedger, cfg hostedGatewayConfig) bool { + if value.Version != 1 || value.ConfigDigest != brokerJSONDigest(cfg) || value.SessionID != cfg.SessionID || + (value.PrincipalDigest != "" && !brokerDigestValid(value.PrincipalDigest)) || + (value.CreateAttempted && value.PrincipalDigest == "") || + (value.SessionCreated && !value.CreateAttempted) || (value.Ready && !value.ExposurePossible) { + return false + } + if !value.ExposurePossible { + return value.Challenge == (hostedChallenge{}) && value.PairID == "" && value.BootstrapDigest == "" && !value.Closed + } + _, nonceValid := hostedCanonicalBytes(value.Challenge.Nonce, 32) + return value.SessionCreated && value.PrincipalDigest != "" && hostedUUIDValid(value.PairID) && + hostedUUIDValid(value.Challenge.BootID) && nonceValid && + brokerDigestValid(value.BootstrapDigest) && value.Challenge.SessionID == cfg.SessionID && + value.Challenge.ConfigurationDigest == brokerJSONDigest(cfg.Image) && + value.Challenge.DeploymentID == cfg.Image.DeploymentID && value.Challenge.Protocol == hostedProtocol && + value.Challenge.AgentName == cfg.Image.Target.AgentName && value.Challenge.AgentVersion == cfg.Image.Target.AgentVersion +} + +func (s *hostedGatewayStore) save(value hostedGatewayLedger) error { + data, err := json.Marshal(value) + if err != nil || len(data) > hostedMaxHandshakeBytes { + return errHostedInvalid + } + name := filepath.Join(s.dir, ".state-"+uuid.NewString()) + file, err := os.OpenFile(name, os.O_CREATE|os.O_EXCL|os.O_WRONLY|syscall.O_NOFOLLOW, 0o600) + if err != nil { + return errHostedInvalid + } + defer os.Remove(name) //nolint:errcheck + if n, err := file.Write(data); err != nil || n != len(data) { + _ = file.Close() + return io.ErrShortWrite + } + if err := file.Sync(); err != nil { + _ = file.Close() + return errHostedInvalid + } + if file.Close() != nil || os.Rename(name, filepath.Join(s.dir, "state.json")) != nil { + return errHostedInvalid + } + directory, err := os.Open(s.dir) + if err != nil { + return errHostedInvalid + } + defer directory.Close() //nolint:errcheck + return directory.Sync() +} + +func (s *hostedGatewayStore) close() { + if s != nil && s.lock != nil { + _ = syscall.Flock(int(s.lock.Fd()), syscall.LOCK_UN) + _ = s.lock.Close() + } +} diff --git a/hosted_transport.go b/hosted_transport.go new file mode 100644 index 0000000..c0dd060 --- /dev/null +++ b/hosted_transport.go @@ -0,0 +1,345 @@ +package main + +import ( + "context" + "errors" + "io" + "log" + "net" + "net/http" + "os" + "sync" + "sync/atomic" + "time" + + "github.com/gorilla/websocket" + "golang.org/x/net/http2" +) + +const ( + hostedWSMaxMessageBytes = 1 << 20 + hostedWSWriteChunkBytes = 256 << 10 + hostedTransportWriteLimit = 30 * time.Second + hostedHTTP2MaxStreams = 64 + hostedHTTP2MaxHeaderBytes = 32 << 10 + hostedHTTP2MaxFrameBytes = 256 << 10 +) + +var errHostedWSMessage = errors.New("hosted transport requires nonempty binary messages of at most 1 MiB") + +// A terminal summary contains only locally assigned role/time and a numeric +// WebSocket close code. Peer close text and transport errors never enter logs. +type hostedChannelObservation struct { + role string + openedAt time.Time + logger *log.Logger + closeCode atomic.Int32 + once sync.Once +} + +func observeHostedChannel(ws *websocket.Conn, role string, openedAt time.Time, logger *log.Logger) *hostedChannelObservation { + if role != "forward" && role != "reverse" { + return nil + } + if logger == nil { + logger = log.Default() + } + o := &hostedChannelObservation{role: role, openedAt: openedAt, logger: logger} + previous := ws.CloseHandler() + ws.SetCloseHandler(func(code int, text string) error { + // Record before the existing reply handler or a lifetime watcher can + // close the pair. Preserve Gorilla's control-frame behavior unchanged. + o.closeCode.CompareAndSwap(0, int32(code)) + return previous(code, text) + }) + return o +} + +func (o *hostedChannelObservation) observeError(err error) { + if o == nil { + return + } + var closed *websocket.CloseError + if errors.As(err, &closed) { + o.closeCode.CompareAndSwap(0, int32(closed.Code)) + } +} + +func (o *hostedChannelObservation) finish() { + if o == nil { + return + } + o.once.Do(func() { + o.logger.Printf("Foundry hosted channel closed role=%s opened_at=%s duration_ms=%d close_code=%d", + o.role, o.openedAt.UTC().Format(time.RFC3339Nano), time.Since(o.openedAt).Milliseconds(), o.closeCode.Load()) + }) +} + +// hostedWSConn carries a byte stream over an already authenticated WebSocket. +// Construction transfers exclusive ownership of the WebSocket to this adapter. +// Whole messages are validated before exposing any bytes to HTTP/2. +type hostedWSConn struct { + ws *websocket.Conn + observation *hostedChannelObservation + done chan struct{} + closeOnce sync.Once + closeErr error + readMu sync.Mutex + readBuf []byte + writeMu sync.Mutex + + deadlineMu sync.Mutex + writeDeadline time.Time + activeWrite *hostedWSWrite +} + +type hostedWSWrite struct { + limit time.Time + timer *time.Timer + generation uint64 + timedOut bool +} + +var _ net.Conn = (*hostedWSConn)(nil) + +func newHostedWSConn(ws *websocket.Conn) *hostedWSConn { + return newHostedObservedWSConn(ws, nil) +} + +func newHostedObservedWSConn(ws *websocket.Conn, observation *hostedChannelObservation) *hostedWSConn { + c := &hostedWSConn{ws: ws, observation: observation, done: make(chan struct{})} + ws.SetReadLimit(hostedWSMaxMessageBytes) + // Remove handshake deadlines. No timer runs while this connection is idle. + // Gorilla's write deadline is only set before handing ownership to callers; + // its setter is not safe concurrently with a writer, nor can it interrupt one. + _ = ws.SetWriteDeadline(time.Time{}) + if err := ws.UnderlyingConn().SetDeadline(time.Time{}); err != nil { + _ = c.Close() + } + return c +} + +func (c *hostedWSConn) Done() <-chan struct{} { return c.done } + +func (c *hostedWSConn) closed() bool { + select { + case <-c.done: + return true + default: + return false + } +} + +func (c *hostedWSConn) Close() error { + c.closeOnce.Do(func() { + close(c.done) + c.deadlineMu.Lock() + if c.activeWrite != nil && c.activeWrite.timer != nil { + c.activeWrite.timer.Stop() + } + c.deadlineMu.Unlock() + c.closeErr = c.ws.Close() + c.observation.finish() + }) + return c.closeErr +} + +func (c *hostedWSConn) Read(p []byte) (int, error) { + c.readMu.Lock() + defer c.readMu.Unlock() + if c.closed() { + return 0, net.ErrClosed + } + if len(p) == 0 { + return 0, nil + } + if len(c.readBuf) == 0 { + kind, reader, err := c.ws.NextReader() + if err != nil { + c.observation.observeError(err) + _ = c.Close() + return 0, hostedTransportError(err) + } + if kind != websocket.BinaryMessage { + _ = c.Close() + return 0, errHostedWSMessage + } + // The extra byte also bounds expanded data if the handshake negotiated + // compression; Gorilla's read limit counts the on-wire message length. + message, err := io.ReadAll(io.LimitReader(reader, hostedWSMaxMessageBytes+1)) + if err != nil { + c.observation.observeError(err) + _ = c.Close() + return 0, hostedTransportError(err) + } + if len(message) == 0 || len(message) > hostedWSMaxMessageBytes { + _ = c.Close() + return 0, errHostedWSMessage + } + c.readBuf = message + } + n := copy(p, c.readBuf) + c.readBuf = c.readBuf[n:] + if len(c.readBuf) == 0 { + c.readBuf = nil + } + return n, nil +} + +func (c *hostedWSConn) Write(p []byte) (int, error) { + c.writeMu.Lock() + defer c.writeMu.Unlock() + if c.closed() { + return 0, net.ErrClosed + } + if len(p) == 0 { + return 0, nil + } + write, err := c.beginWrite() + if err != nil { + _ = c.Close() + return 0, err + } + n := 0 + for n < len(p) { + end := n + min(len(p)-n, hostedWSWriteChunkBytes) + if err = c.ws.WriteMessage(websocket.BinaryMessage, p[n:end]); err != nil { + break + } + // A failed message may have an incomplete frame on the wire. Only + // completed messages count toward the accepted prefix; never resend it. + n = end + } + if c.finishWrite(write) { + err = os.ErrDeadlineExceeded + } + if err != nil { + c.observation.observeError(err) + _ = c.Close() + } + return n, hostedTransportError(err) +} + +func hostedTransportError(err error) error { + var networkError net.Error + if errors.As(err, &networkError) && networkError.Timeout() { + // Gorilla deliberately removes the wrapped syscall error. Restore + // net.Conn's deadline sentinel while keeping this connection terminal. + return os.ErrDeadlineExceeded + } + return err +} + +func (c *hostedWSConn) beginWrite() (*hostedWSWrite, error) { + c.deadlineMu.Lock() + defer c.deadlineMu.Unlock() + if c.closed() { + return nil, net.ErrClosed + } + now := time.Now() + if !c.writeDeadline.IsZero() && !c.writeDeadline.After(now) { + return nil, os.ErrDeadlineExceeded + } + write := &hostedWSWrite{limit: now.Add(hostedTransportWriteLimit)} + c.activeWrite = write + c.armWriteTimer(write) + return write, nil +} + +// armWriteTimer runs under deadlineMu. A generation prevents an already queued +// timer callback from applying a deadline that a caller extended or cleared. +func (c *hostedWSConn) armWriteTimer(write *hostedWSWrite) { + if write.timer != nil { + write.timer.Stop() + } + deadline := write.limit + if !c.writeDeadline.IsZero() && c.writeDeadline.Before(deadline) { + deadline = c.writeDeadline + } + write.generation++ + generation := write.generation + write.timer = time.AfterFunc(time.Until(deadline), func() { + c.deadlineMu.Lock() + if c.activeWrite != write || write.generation != generation || c.closed() { + c.deadlineMu.Unlock() + return + } + write.timedOut = true + c.deadlineMu.Unlock() + _ = c.Close() + }) +} + +func (c *hostedWSConn) finishWrite(write *hostedWSWrite) bool { + c.deadlineMu.Lock() + defer c.deadlineMu.Unlock() + write.timer.Stop() + c.activeWrite = nil + return write.timedOut +} + +func (c *hostedWSConn) LocalAddr() net.Addr { return c.ws.LocalAddr() } +func (c *hostedWSConn) RemoteAddr() net.Addr { return c.ws.RemoteAddr() } + +func (c *hostedWSConn) SetDeadline(deadline time.Time) error { + if err := c.SetReadDeadline(deadline); err != nil { + return err + } + return c.SetWriteDeadline(deadline) +} + +func (c *hostedWSConn) SetReadDeadline(deadline time.Time) error { + return c.ws.UnderlyingConn().SetReadDeadline(deadline) +} + +func (c *hostedWSConn) SetWriteDeadline(deadline time.Time) error { + c.deadlineMu.Lock() + defer c.deadlineMu.Unlock() + if c.closed() { + return net.ErrClosed + } + c.writeDeadline = deadline + if c.activeWrite != nil && !c.activeWrite.timedOut { + c.armWriteTimer(c.activeWrite) + } + return nil +} + +// Call ClientConn.RoundTrip directly. A Transport.RoundTrip would add connection +// selection and transparent retries, which cannot preserve mutation ownership. +func newHostedHTTP2ClientConn(conn net.Conn) (*http2.ClientConn, error) { + // ConfigureTransports initializes x/net's underlying implementation on + // both Go 1.26 and 1.27. Neither transport's RoundTrip is ever used. + transport, err := http2.ConfigureTransports(&http.Transport{}) + if err != nil { + _ = conn.Close() + return nil, err + } + transport.AllowHTTP = true + transport.DisableCompression = true + transport.StrictMaxConcurrentStreams = true + transport.MaxHeaderListSize = hostedHTTP2MaxHeaderBytes + transport.MaxReadFrameSize = hostedHTTP2MaxFrameBytes + transport.WriteByteTimeout = hostedTransportWriteLimit + client, err := transport.NewClientConn(conn) + if err != nil { + _ = conn.Close() + } + return client, err +} + +func serveHostedHTTP2(ctx context.Context, conn net.Conn, handler http.Handler) { + stop := context.AfterFunc(ctx, func() { _ = conn.Close() }) + defer stop() + defer conn.Close() + server := &http2.Server{ + MaxConcurrentStreams: hostedHTTP2MaxStreams, + MaxReadFrameSize: hostedHTTP2MaxFrameBytes, + WriteByteTimeout: hostedTransportWriteLimit, + } + server.ServeConn(conn, &http2.ServeConnOpts{ + Context: ctx, + BaseConfig: &http.Server{MaxHeaderBytes: hostedHTTP2MaxHeaderBytes}, + Handler: handler, + }) +} diff --git a/hosted_transport_test.go b/hosted_transport_test.go new file mode 100644 index 0000000..0e31217 --- /dev/null +++ b/hosted_transport_test.go @@ -0,0 +1,694 @@ +package main + +import ( + "bufio" + "bytes" + "context" + "errors" + "io" + "net" + "net/http" + "net/http/httptest" + "os" + "strings" + "sync" + "sync/atomic" + "testing" + "time" + + "github.com/gorilla/websocket" + "golang.org/x/net/http2" +) + +func hostedTestWSPair(t *testing.T, compression bool, gate *hostedTestWriteGate) (*websocket.Conn, *websocket.Conn) { + t.Helper() + accepted := make(chan *websocket.Conn, 1) + upgradeError := make(chan error, 1) + upgrader := websocket.Upgrader{WriteBufferSize: 1024, EnableCompression: compression} + server := httptest.NewServer(http.HandlerFunc(func(w http.ResponseWriter, r *http.Request) { + conn, err := upgrader.Upgrade(w, r, nil) + if err != nil { + upgradeError <- err + return + } + accepted <- conn + })) + t.Cleanup(server.Close) + dialer := websocket.Dialer{ + WriteBufferSize: hostedWSWriteChunkBytes, + EnableCompression: compression, + HandshakeTimeout: 5 * time.Second, + } + if gate != nil { + dialer.NetDialContext = func(ctx context.Context, network, address string) (net.Conn, error) { + conn, err := (&net.Dialer{}).DialContext(ctx, network, address) + if err != nil { + return nil, err + } + gate.Conn = conn + return gate, nil + } + } + client, response, err := dialer.Dial("ws"+strings.TrimPrefix(server.URL, "http"), nil) + if response != nil && response.Body != nil { + _ = response.Body.Close() + } + if err != nil { + t.Fatal("localhost WebSocket handshake failed") + } + var peer *websocket.Conn + select { + case peer = <-accepted: + case <-upgradeError: + t.Fatal("localhost WebSocket upgrade failed") + case <-time.After(5 * time.Second): + t.Fatal("localhost WebSocket upgrade did not complete") + } + t.Cleanup(func() { + _ = client.Close() + _ = peer.Close() + }) + return client, peer +} + +func hostedTestDone(t *testing.T, done <-chan struct{}) { + t.Helper() + select { + case <-done: + case <-time.After(5 * time.Second): + t.Fatal("transport operation did not finish") + } +} + +func TestHostedWSConnReadsValidatedMessages(t *testing.T) { + client, peer := hostedTestWSPair(t, false, nil) + conn := newHostedWSConn(client) + defer conn.Close() + if conn.LocalAddr() == nil || conn.RemoteAddr() == nil { + t.Fatal("transport lost its endpoint addresses") + } + payload := bytes.Repeat([]byte{0x41}, hostedWSMaxMessageBytes) + sent := make(chan error, 1) + go func() { + writer, err := peer.NextWriter(websocket.BinaryMessage) + if err == nil { + // Several writes and a small peer buffer force WebSocket continuation + // frames; the receiver still sees a single bounded message. + for offset := 0; offset < len(payload) && err == nil; { + end := min(offset+719, len(payload)) + _, err = writer.Write(payload[offset:end]) + offset = end + } + if err == nil { + err = writer.Close() + } + } + if err == nil { + err = peer.WriteMessage(websocket.BinaryMessage, []byte("tail")) + } + sent <- err + }() + if err := conn.SetReadDeadline(time.Now().Add(5 * time.Second)); err != nil { + t.Fatal("could not bound test read") + } + got := make([]byte, len(payload)+4) + for offset := 0; offset < len(got); { + end := min(offset+113, len(got)) + n, err := conn.Read(got[offset:end]) + if err != nil || n == 0 { + t.Fatal("valid fragmented stream was truncated") + } + offset += n + } + if !bytes.Equal(got[:len(payload)], payload) || string(got[len(payload):]) != "tail" || <-sent != nil { + t.Fatal("message boundaries changed the byte stream") + } +} + +func TestHostedWSConnRejectsInvalidMessagesWithoutExposingBytes(t *testing.T) { + for _, test := range []struct { + name string + kind int + length int + compression bool + incomplete bool + }{ + {name: "empty", kind: websocket.BinaryMessage}, + {name: "text", kind: websocket.TextMessage, length: 4}, + {name: "oversized fragmented", kind: websocket.BinaryMessage, length: hostedWSMaxMessageBytes + 1}, + {name: "oversized expanded", kind: websocket.BinaryMessage, length: hostedWSMaxMessageBytes + 1, compression: true}, + {name: "incomplete fragmented", kind: websocket.BinaryMessage, length: 4096, incomplete: true}, + } { + t.Run(test.name, func(t *testing.T) { + client, peer := hostedTestWSPair(t, test.compression, nil) + conn := newHostedWSConn(client) + defer conn.Close() + _ = conn.SetReadDeadline(time.Now().Add(5 * time.Second)) + sent := make(chan struct{}) + go func() { + defer close(sent) + writer, err := peer.NextWriter(test.kind) + if err == nil { + _, err = writer.Write(bytes.Repeat([]byte{'x'}, test.length)) + if err == nil && !test.incomplete { + _ = writer.Close() + } + } + if test.incomplete { + _ = peer.Close() + } + }() + if n, err := conn.Read(make([]byte, 31)); n != 0 || err == nil { + t.Fatal("invalid or incomplete message exposed stream bytes") + } + hostedTestDone(t, conn.Done()) + hostedTestDone(t, sent) + // A caller that continues after failure must not re-enter Gorilla's + // failed-read path, which eventually panics on repeated reads. + for range 1100 { + if n, err := conn.Read(make([]byte, 1)); n != 0 || !errors.Is(err, net.ErrClosed) { + t.Fatal("failed connection became readable again") + } + } + }) + } +} + +func TestHostedWSConnWritesBoundedMessages(t *testing.T) { + client, peer := hostedTestWSPair(t, false, nil) + conn := newHostedWSConn(client) + defer conn.Close() + payload := bytes.Repeat([]byte{0x42}, 2*hostedWSWriteChunkBytes+17) + if n, err := conn.Write(nil); n != 0 || err != nil { + t.Fatal("zero-length write failed") + } + writeResult := make(chan error, 1) + go func() { + n, err := conn.Write(payload) + if err == nil && n != len(payload) { + err = io.ErrShortWrite + } + writeResult <- err + }() + _ = peer.SetReadDeadline(time.Now().Add(5 * time.Second)) + var got []byte + for _, size := range []int{hostedWSWriteChunkBytes, hostedWSWriteChunkBytes, 17} { + kind, part, err := peer.ReadMessage() + if err != nil || kind != websocket.BinaryMessage || len(part) != size { + t.Fatal("write emitted an empty, non-binary or incorrectly sized message") + } + got = append(got, part...) + } + if <-writeResult != nil || !bytes.Equal(got, payload) { + t.Fatal("chunked write changed the stream") + } +} + +func TestHostedWSConnConcurrentWritesPreserveCompletePrefixes(t *testing.T) { + client, peer := hostedTestWSPair(t, false, nil) + conn := newHostedWSConn(client) + defer conn.Close() + const count = 4 + written := make(chan error, count) + for id := range count { + go func() { + payload := bytes.Repeat([]byte{byte(id + 1)}, hostedWSWriteChunkBytes+17) + n, err := conn.Write(payload) + if err == nil && n != len(payload) { + err = io.ErrShortWrite + } + written <- err + }() + } + _ = peer.SetReadDeadline(time.Now().Add(5 * time.Second)) + seen := make(map[byte]bool) + for range count { + kind, first, err := peer.ReadMessage() + if err != nil || kind != websocket.BinaryMessage || len(first) != hostedWSWriteChunkBytes { + t.Fatal("concurrent write lost its first chunk") + } + kind, last, err := peer.ReadMessage() + id := first[0] + if err != nil || kind != websocket.BinaryMessage || len(last) != 17 || seen[id] || + !bytes.Equal(first, bytes.Repeat([]byte{id}, len(first))) || + !bytes.Equal(last, bytes.Repeat([]byte{id}, len(last))) { + t.Fatal("concurrent writes interleaved or duplicated stream bytes") + } + seen[id] = true + } + for range count { + if <-written != nil { + t.Fatal("concurrent write failed") + } + } +} + +var errHostedTestWrite = errors.New("injected transport write failure") + +// hostedTestWriteGate controls an actual TCP connection after its WS handshake. +// It makes write failure and blocked-write cancellation independent of TCP buffer +// capacity and scheduling. No transport implementation is replaced. +type hostedTestWriteGate struct { + net.Conn + mu sync.Mutex + enabled bool + block bool + failAfter int + calls int + entered chan struct{} + closed chan struct{} + enterOnce sync.Once + closeOnce sync.Once +} + +func newHostedTestWriteGate() *hostedTestWriteGate { + return &hostedTestWriteGate{entered: make(chan struct{}), closed: make(chan struct{})} +} + +func (g *hostedTestWriteGate) arm(block bool, failAfter int) { + g.mu.Lock() + defer g.mu.Unlock() + g.enabled, g.block, g.failAfter = true, block, failAfter +} + +func (g *hostedTestWriteGate) Write(p []byte) (int, error) { + g.mu.Lock() + enabled, block := g.enabled, g.block + if enabled { + g.calls++ + } + fail := enabled && !block && g.calls > g.failAfter + g.mu.Unlock() + if enabled { + g.enterOnce.Do(func() { close(g.entered) }) + } + if block { + <-g.closed + return 0, net.ErrClosed + } + if fail { + n, _ := g.Conn.Write(p[:min(17, len(p))]) + return n, errHostedTestWrite + } + return g.Conn.Write(p) +} + +func (g *hostedTestWriteGate) Close() error { + g.closeOnce.Do(func() { close(g.closed) }) + return g.Conn.Close() +} + +func TestHostedWSConnPartialWriteCannotReplay(t *testing.T) { + gate := newHostedTestWriteGate() + client, peer := hostedTestWSPair(t, false, gate) + conn := newHostedWSConn(client) + reader := newHostedWSConn(peer) + defer conn.Close() + defer reader.Close() + gate.arm(false, 1) + payload := bytes.Repeat([]byte{0x43}, 2*hostedWSWriteChunkBytes) + type result struct { + n int + err error + } + written := make(chan result, 1) + go func() { + n, err := conn.Write(payload) + written <- result{n, err} + }() + _ = reader.SetReadDeadline(time.Now().Add(5 * time.Second)) + got := make([]byte, hostedWSWriteChunkBytes) + if _, err := io.ReadFull(reader, got); err != nil || !bytes.Equal(got, payload[:len(got)]) { + t.Fatal("completed write prefix was lost") + } + if n, err := reader.Read(make([]byte, 31)); n != 0 || err == nil { + t.Fatal("partial failed message exposed bytes") + } + resultValue := <-written + if resultValue.n != hostedWSWriteChunkBytes || !errors.Is(resultValue.err, errHostedTestWrite) { + t.Fatal("partial write did not report its completed prefix and error") + } + hostedTestDone(t, conn.Done()) + if n, err := conn.Write(payload); n != 0 || !errors.Is(err, net.ErrClosed) { + t.Fatal("failed connection allowed another write") + } + gate.mu.Lock() + calls := gate.calls + gate.mu.Unlock() + if calls != 2 { + t.Fatal("failed write was retried") + } +} + +func TestHostedWSConnActiveWriteDeadlineCanChange(t *testing.T) { + gate := newHostedTestWriteGate() + client, _ := hostedTestWSPair(t, false, gate) + conn := newHostedWSConn(client) + defer conn.Close() + gate.arm(true, 0) + written := make(chan error, 1) + started := time.Now() + go func() { + n, err := conn.Write([]byte("blocked")) + if n != 0 { + err = errors.New("blocked write reported accepted bytes") + } + written <- err + }() + hostedTestDone(t, gate.entered) + conn.deadlineMu.Lock() + limit := conn.activeWrite.limit + conn.deadlineMu.Unlock() + if limit.Before(started.Add(hostedTransportWriteLimit)) || limit.After(time.Now().Add(hostedTransportWriteLimit)) { + t.Fatal("active write lacks its default finite limit") + } + if err := conn.SetWriteDeadline(time.Now().Add(50 * time.Millisecond)); err != nil { + t.Fatal("could not shorten active deadline") + } + if err := conn.SetWriteDeadline(time.Time{}); err != nil { + t.Fatal("could not clear caller deadline") + } + select { + case <-conn.Done(): + t.Fatal("superseded deadline closed the active write") + case <-time.After(100 * time.Millisecond): + } + if err := conn.SetWriteDeadline(time.Now().Add(-time.Second)); err != nil { + t.Fatal("could not expire pending write") + } + select { + case err := <-written: + if !errors.Is(err, os.ErrDeadlineExceeded) { + t.Fatal("blocked write did not report deadline expiry") + } + case <-time.After(5 * time.Second): + t.Fatal("changing the deadline did not interrupt the writer") + } + hostedTestDone(t, conn.Done()) +} + +func TestHostedWSConnExpiredIdleDeadlineDoesNotCloseUntilIO(t *testing.T) { + client, _ := hostedTestWSPair(t, false, nil) + conn := newHostedWSConn(client) + defer conn.Close() + if err := conn.SetDeadline(time.Now().Add(-time.Second)); err != nil { + t.Fatal("could not set idle deadline") + } + conn.deadlineMu.Lock() + active := conn.activeWrite + conn.deadlineMu.Unlock() + if active != nil || conn.closed() { + t.Fatal("idle connection acquired a write timer or closed without I/O") + } + if n, err := conn.Write([]byte("expired")); n != 0 || !errors.Is(err, os.ErrDeadlineExceeded) { + t.Fatal("past write deadline allowed bytes") + } + hostedTestDone(t, conn.Done()) +} + +func TestHostedWSConnReadDeadlineAndConcurrentClose(t *testing.T) { + client, _ := hostedTestWSPair(t, false, nil) + conn := newHostedWSConn(client) + read := make(chan error, 1) + go func() { + _, err := conn.Read(make([]byte, 1)) + read <- err + }() + if err := conn.SetReadDeadline(time.Now().Add(-time.Second)); err != nil { + t.Fatal("could not expire pending read") + } + select { + case err := <-read: + if !errors.Is(err, os.ErrDeadlineExceeded) { + t.Fatal("blocked read did not report deadline expiry") + } + case <-time.After(5 * time.Second): + t.Fatal("read deadline did not unblock the connection") + } + var group sync.WaitGroup + for range 16 { + group.Go(func() { _ = conn.Close() }) + } + group.Wait() + hostedTestDone(t, conn.Done()) + if !errors.Is(conn.SetWriteDeadline(time.Time{}), net.ErrClosed) { + t.Fatal("closed connection accepted a deadline") + } +} + +type hostedTestHTTP2 struct { + client *http2.ClientConn + clientConn *hostedWSConn + serverConn *hostedWSConn + cancel context.CancelFunc + done <-chan struct{} +} + +func hostedTestHTTP2Pair(t *testing.T, handler http.Handler) *hostedTestHTTP2 { + t.Helper() + clientWS, serverWS := hostedTestWSPair(t, false, nil) + clientConn, serverConn := newHostedWSConn(clientWS), newHostedWSConn(serverWS) + ctx, cancel := context.WithCancel(context.Background()) + done := make(chan struct{}) + go func() { + defer close(done) + serveHostedHTTP2(ctx, serverConn, handler) + }() + client, err := newHostedHTTP2ClientConn(clientConn) + if err != nil { + cancel() + t.Fatal("HTTP/2 client construction failed") + } + t.Cleanup(func() { + cancel() + _ = client.Close() + _ = clientConn.Close() + _ = serverConn.Close() + hostedTestDone(t, done) + }) + return &hostedTestHTTP2{client, clientConn, serverConn, cancel, done} +} + +func TestHostedHTTP2NDJSONStreamingAndCancellation(t *testing.T) { + cancelled := make(chan struct{}) + pair := hostedTestHTTP2Pair(t, http.HandlerFunc(func(w http.ResponseWriter, r *http.Request) { + if r.URL.Path == "/health" { + w.WriteHeader(http.StatusNoContent) + return + } + w.Header().Set("Content-Type", "application/x-ndjson") + _, _ = io.WriteString(w, "{\"sequence\":1}\n") + _ = http.NewResponseController(w).Flush() + <-r.Context().Done() + close(cancelled) + })) + ctx, cancel := context.WithTimeout(context.Background(), 5*time.Second) + defer cancel() + request, _ := http.NewRequestWithContext(ctx, http.MethodGet, "http://hosted/stream", nil) + response, err := pair.client.RoundTrip(request) + if err != nil { + t.Fatal("stream request failed") + } + defer response.Body.Close() + line, err := bufio.NewReader(response.Body).ReadString('\n') + if err != nil || line != "{\"sequence\":1}\n" { + t.Fatal("NDJSON was not available while the handler was active") + } + select { + case <-cancelled: + t.Fatal("stream ended before client cancellation") + default: + } + cancel() + hostedTestDone(t, cancelled) + freshContext, freshCancel := context.WithTimeout(context.Background(), 5*time.Second) + defer freshCancel() + request, _ = http.NewRequestWithContext(freshContext, http.MethodGet, "http://hosted/health", nil) + response, err = pair.client.RoundTrip(request) + if err != nil { + t.Fatal("cancelling a stream destroyed the shared HTTP/2 connection") + } + _ = response.Body.Close() + if response.StatusCode != http.StatusNoContent || pair.clientConn.closed() || pair.serverConn.closed() { + t.Fatal("connection was not reusable after stream cancellation") + } +} + +func TestHostedHTTP2ConcurrentRequests(t *testing.T) { + const count = 12 + entered := make(chan struct{}, count) + release := make(chan struct{}) + pair := hostedTestHTTP2Pair(t, http.HandlerFunc(func(w http.ResponseWriter, r *http.Request) { + entered <- struct{}{} + select { + case <-release: + _, _ = io.Copy(w, r.Body) + case <-r.Context().Done(): + } + })) + ctx, cancel := context.WithTimeout(context.Background(), 5*time.Second) + defer cancel() + results := make(chan error, count) + for id := range count { + go func() { + payload := bytes.Repeat([]byte{byte(id + 1)}, 64<<10) + request, _ := http.NewRequestWithContext(ctx, http.MethodPost, "http://hosted/echo", bytes.NewReader(payload)) + response, err := pair.client.RoundTrip(request) + if err == nil { + var got []byte + got, err = io.ReadAll(response.Body) + _ = response.Body.Close() + if err == nil && !bytes.Equal(got, payload) { + err = errors.New("multiplexed request content changed") + } + } + results <- err + }() + } + for range count { + select { + case <-entered: + case <-ctx.Done(): + t.Fatal("requests were serialized instead of multiplexed") + } + } + close(release) + for range count { + if <-results != nil { + t.Fatal("concurrent HTTP/2 request failed") + } + } +} + +func TestHostedHTTP2HeaderAndStreamBounds(t *testing.T) { + var called atomic.Int32 + pair := hostedTestHTTP2Pair(t, http.HandlerFunc(func(w http.ResponseWriter, r *http.Request) { + called.Add(1) + if r.URL.Path == "/large-response" { + w.Header().Set("X-Large", strings.Repeat("x", 2*hostedHTTP2MaxHeaderBytes)) + } + w.WriteHeader(http.StatusOK) + })) + ctx, cancel := context.WithTimeout(context.Background(), 5*time.Second) + defer cancel() + request, _ := http.NewRequestWithContext(ctx, http.MethodGet, "http://hosted/ready", nil) + response, err := pair.client.RoundTrip(request) + if err != nil { + t.Fatal("initial settings exchange failed") + } + _ = response.Body.Close() + if pair.client.State().MaxConcurrentStreams != hostedHTTP2MaxStreams { + t.Fatal("server did not advertise the stream bound") + } + request, _ = http.NewRequestWithContext(ctx, http.MethodGet, "http://hosted/large-request", nil) + request.Header.Set("X-Large", strings.Repeat("x", 2*hostedHTTP2MaxHeaderBytes)) + if response, err = pair.client.RoundTrip(request); err == nil { + _ = response.Body.Close() + t.Fatal("oversized request headers were accepted") + } + if called.Load() != 1 { + t.Fatal("oversized headers reached the handler") + } + request, _ = http.NewRequestWithContext(ctx, http.MethodGet, "http://hosted/large-response", nil) + if response, err = pair.client.RoundTrip(request); err == nil { + _ = response.Body.Close() + t.Fatal("oversized response headers were accepted") + } +} + +func TestHostedHTTP2SaturatedStreamCanBeCancelledWithoutDispatch(t *testing.T) { + entered := make(chan struct{}, hostedHTTP2MaxStreams+1) + release := make(chan struct{}) + pair := hostedTestHTTP2Pair(t, http.HandlerFunc(func(w http.ResponseWriter, r *http.Request) { + entered <- struct{}{} + select { + case <-release: + w.WriteHeader(http.StatusNoContent) + case <-r.Context().Done(): + } + })) + ctx, cancel := context.WithTimeout(context.Background(), 5*time.Second) + defer cancel() + results := make(chan error, hostedHTTP2MaxStreams) + for range hostedHTTP2MaxStreams { + go func() { + request, _ := http.NewRequestWithContext(ctx, http.MethodGet, "http://hosted/hold", nil) + response, err := pair.client.RoundTrip(request) + if response != nil { + _ = response.Body.Close() + } + results <- err + }() + } + for range hostedHTTP2MaxStreams { + select { + case <-entered: + case <-ctx.Done(): + t.Fatal("stream capacity could not be filled") + } + } + queued, cancelQueued := context.WithTimeout(ctx, 100*time.Millisecond) + defer cancelQueued() + request, _ := http.NewRequestWithContext(queued, http.MethodGet, "http://hosted/queued", nil) + response, err := pair.client.RoundTrip(request) + if response != nil { + _ = response.Body.Close() + } + if !errors.Is(err, context.DeadlineExceeded) { + t.Fatal("queued stream did not observe cancellation") + } + select { + case <-entered: + t.Fatal("request exceeded the concurrent stream bound") + default: + } + close(release) + for range hostedHTTP2MaxStreams { + if <-results != nil { + t.Fatal("queued cancellation disturbed an admitted stream") + } + } +} + +func TestHostedHTTP2AcceptedRequestIsNotReplayedAfterDisconnect(t *testing.T) { + var calls, replays atomic.Int32 + sever := make(chan *hostedWSConn, 1) + pair := hostedTestHTTP2Pair(t, http.HandlerFunc(func(w http.ResponseWriter, r *http.Request) { + _, _ = io.Copy(io.Discard, r.Body) + calls.Add(1) + _ = (<-sever).Close() + })) + sever <- pair.serverConn + ctx, cancel := context.WithTimeout(context.Background(), 5*time.Second) + defer cancel() + request, _ := http.NewRequestWithContext(ctx, http.MethodPost, "http://hosted/mutation", strings.NewReader("once")) + request.Header.Set("Idempotency-Key", "fixture") + request.GetBody = func() (io.ReadCloser, error) { + replays.Add(1) + return io.NopCloser(strings.NewReader("once")), nil + } + if response, err := pair.client.RoundTrip(request); err == nil { + _ = response.Body.Close() + t.Fatal("severed accepted request unexpectedly succeeded") + } + hostedTestDone(t, pair.clientConn.Done()) + if calls.Load() != 1 || replays.Load() != 0 { + t.Fatal("accepted mutation was replayed after connection loss") + } + request, _ = http.NewRequestWithContext(ctx, http.MethodPost, "http://hosted/after-loss", nil) + if response, err := pair.client.RoundTrip(request); err == nil { + _ = response.Body.Close() + t.Fatal("closed ClientConn acquired a replacement connection") + } + if calls.Load() != 1 { + t.Fatal("closed connection dispatched another mutation") + } +} + +func TestHostedHTTP2ServerContextClosesIdleConnection(t *testing.T) { + pair := hostedTestHTTP2Pair(t, http.HandlerFunc(func(http.ResponseWriter, *http.Request) {})) + pair.cancel() + hostedTestDone(t, pair.done) + hostedTestDone(t, pair.serverConn.Done()) + hostedTestDone(t, pair.clientConn.Done()) +} diff --git a/main.go b/main.go index e740f62..fb2e376 100644 --- a/main.go +++ b/main.go @@ -3,10 +3,29 @@ package main import ( "log" "net/http" + "os" "time" ) func main() { + if handled, err := maybeServeHosted(os.Args[1:]); handled { + if err != nil { + log.Fatal("Foundry hosted lifetime unavailable; inspect the ownership ledger before replacement") + } + return + } + if handled, err := maybeServeBroker(os.Args[1:]); handled { + if err != nil { + log.Fatal("Foundry lifecycle broker failed") + } + return + } + if handled, err := maybeServeACP(os.Args[1:], os.Stdin, os.Stdout); handled { + if err != nil { + log.Fatal("Foundry ACP bridge failed") + } + return + } cfg := loadConfig() if err := cfg.validate(); err != nil { log.Fatal(err) diff --git a/remote_header_preflight_test.go b/remote_header_preflight_test.go new file mode 100644 index 0000000..62f35e2 --- /dev/null +++ b/remote_header_preflight_test.go @@ -0,0 +1,136 @@ +package main + +import ( + "context" + "encoding/json" + "net/http" + "net/http/httptest" + "os" + "path/filepath" + "sync/atomic" + "testing" +) + +func invalidRemoteHeaderTokens() map[string]func(string) string { + return map[string]func(string) string{ + "signature-nul": func(token string) string { return token + "\x00" }, + "signature-lf": func(token string) string { return token + "\n" }, + "signature-cr": func(token string) string { return token + "\r" }, + "header-del": func(token string) string { return "\x7f" + token }, + } +} + +func TestBrokerInvalidRemoteHeaderDoesNotReserveSubmission(t *testing.T) { + for _, phase := range []string{"create", "inference"} { + for name, corrupt := range invalidRemoteHeaderTokens() { + t.Run(phase+"/"+name, func(t *testing.T) { + f := newBrokerFixture(t, "success") + cfg := brokerTestConfig(t, f) + c := brokerTestContext(cfg) + failAt, expectedPosts := int64(3), int64(0) + if phase == "inference" { + failAt, expectedPosts = 5, 1 + } + var tokens, posts atomic.Int64 + provider := brokerEvidenceTokenProvider(func(context.Context) (string, error) { + if tokens.Add(1) == failAt { + return corrupt(brokerTestToken()), nil + } + return brokerTestToken(), nil + }) + client := newBrokerHTTPClient() + transport := client.Transport + t.Cleanup(transport.(*http.Transport).CloseIdleConnections) + client.Transport = brokerFixtureTransport(func(request *http.Request) (*http.Response, error) { + if request.Method == http.MethodPost { + posts.Add(1) + } + return transport.RoundTrip(request) + }) + b, err := newLifecycleBroker(t.Context(), cfg, provider, client) + if err != nil { + t.Fatal("could not initialize header preflight fixture") + } + server := httptest.NewServer(b) + t.Cleanup(func() { b.close(); server.Close() }) + status, _, err := brokerTestHTTP(t.Context(), server.URL, brokerResponsesPath, c, brokerTestBody("")) + if err != nil || status == http.StatusOK || tokens.Load() < failAt { + t.Fatal("fixture did not reach the invalid header boundary") + } + if posts.Load() != expectedPosts { + t.Error("invalid authorization reached transport submission") + } + b.close() + server.Close() + raw, err := os.ReadFile(filepath.Join(cfg.stateDir, "state.json")) + var ledger brokerLedger + if err != nil || json.Unmarshal(raw, &ledger) != nil || !brokerLedgerValid(&ledger, cfg.configDigest) { + t.Fatal("could not inspect durable header-failure ownership") + } + owner := ledger.Sessions[brokerJSONDigest(c.Owner)] + if owner == nil || owner.CreateState == "intent" || + (phase == "create" && (owner.CreateState != "none" || owner.RemoteID != "")) { + t.Fatal("invalid authorization stranded an unsent creation") + } + invocation := owner.Prompts[c.promptKey()].Invocations[c.InvocationSequence] + if invocation.State == "intent" || invocation.State == "uncertain" { + t.Fatal("invalid authorization stranded an unsent inference") + } + _, restarted := startBrokerTest(t, cfg) + proof := brokerTestControl(t, restarted.URL, brokerSettlePath, c) + if !proof.SettlementProven || proof.CreatePending || proof.AmbiguousInvocations != 0 { + t.Fatal("restart could not settle definitely-unsent authorization failure") + } + proof = brokerTestControl(t, restarted.URL, brokerRetirePath, c) + creates, inferences, stops, deletes := f.counts() + if !proof.RetirementProven || int64(creates) != expectedPosts || inferences != 0 || stops != 0 || + int64(deletes) != expectedPosts { + t.Fatal("header-failure recovery replayed work or lost owned retirement") + } + }) + } + } +} + +func TestHostedInvalidRemoteHeaderDoesNotReserveCreation(t *testing.T) { + for name, corrupt := range invalidRemoteHeaderTokens() { + t.Run(name, func(t *testing.T) { + f := newHostedGatewayTestFixture(t, hostedGatewayTestOptions{}) + var posts atomic.Int64 + transport := f.gateway.httpClient.Transport + f.gateway.httpClient.Transport = hostedGatewayTestRoundTripper(func(request *http.Request) (*http.Response, error) { + if request.Method == http.MethodPost { + posts.Add(1) + } + return transport.RoundTrip(request) + }) + f.gateway.provider = hostedGatewayTestTokenProvider(func(context.Context) (string, error) { + token := f.token.Load().(string) + if f.tokenCalls.Add(1) == 4 { + return corrupt(token), nil + } + return token, nil + }) + if f.initialize() == nil || f.tokenCalls.Load() != 4 { + t.Fatal("fixture did not reach the invalid creation header") + } + if posts.Load() != 0 || f.creates.Load() != 0 || f.dials.Load() != 0 || f.httpCalls.Load() != 3 { + t.Error("invalid creation header reached transport or channel setup") + } + f.assertNoBootstrap(t) + store, ledger, err := openHostedGatewayStore(f.settings.stateDir, f.settings.config) + if err != nil { + t.Fatal("could not reopen gateway after local header rejection") + } + defer store.close() + if ledger.CreateAttempted || ledger.SessionCreated || ledger.ExposurePossible { + t.Fatal("local header rejection stranded the hosted creation") + } + restarted := &hostedGateway{cfg: f.settings.config, provider: f.gateway.provider, + httpClient: f.gateway.httpClient, store: store, ledger: ledger} + if restarted.ensureSession(t.Context()) != nil || posts.Load() != 1 || f.creates.Load() != 1 || !restarted.ledger.SessionCreated { + t.Fatal("definitely-unsent header rejection prevented a later owned creation") + } + }) + } +} diff --git a/store_directory.go b/store_directory.go new file mode 100644 index 0000000..0984e88 --- /dev/null +++ b/store_directory.go @@ -0,0 +1,112 @@ +package main + +import ( + "os" + "path/filepath" + "syscall" +) + +// The permanent lock is also an initialization witness. Only its exclusive +// creator may initialize absent state; existing ledgers do not need a new marker. +func openStoreLock(dir, name string, syncParent func(string) error) (*os.File, bool, error) { + if err := makeStoreDirectory(dir, syncParent); err != nil { + return nil, false, err + } + info, err := os.Lstat(dir) + if err != nil { + return nil, false, err + } + if stat, ok := info.Sys().(*syscall.Stat_t); !ok || !info.IsDir() || info.Mode().Perm() != 0o700 || int(stat.Uid) != os.Geteuid() { + return nil, false, os.ErrPermission + } + path := filepath.Join(dir, name) + lock, err := os.OpenFile(path, os.O_CREATE|os.O_EXCL|os.O_RDWR|syscall.O_NOFOLLOW, 0o600) + created := err == nil + if os.IsExist(err) { + lock, err = os.OpenFile(path, os.O_RDWR|syscall.O_NOFOLLOW, 0o600) + } + if err != nil { + return nil, false, err + } + if info, err := lock.Stat(); err != nil || !brokerPrivateFile(info) { + _ = lock.Close() + return nil, false, os.ErrPermission + } + if err := flockStoreFile(lock, created); err != nil { + _ = lock.Close() + return nil, false, err + } + return lock, created, nil +} + +func flockStoreFile(lock *os.File, created bool) error { + operation := syscall.LOCK_EX | syscall.LOCK_NB + if created { + _, err := os.Lstat(filepath.Join(filepath.Dir(lock.Name()), "state.json")) + if err != nil && !os.IsNotExist(err) { + return err + } + if os.IsNotExist(err) { + // Keep the creator's O_EXCL authority on the same descriptor + // while an earlier locker rejects the absent ledger. With an + // existing legacy ledger, that locker may be its active owner, + // so recovery must retain nonblocking exclusion instead. + operation = syscall.LOCK_EX + } + } + for { + err := syscall.Flock(int(lock.Fd()), operation) + if err != syscall.EINTR { + return err + } + } +} + +// Persist a directory entry before descending into it or creating ownership +// records. An existing directory may remain after a failed sync or be created +// by a competing initializer, so it needs the same parent barrier. Recursion +// stops at the first existing ancestor; every new descendant follows its sync. +func makeStoreDirectory(dir string, syncParent func(string) error) error { + parent := filepath.Dir(dir) + info, err := os.Stat(dir) + if err == nil { + if info.IsDir() { + if parent == dir { + return nil + } + return syncParent(parent) + } + return &os.PathError{Op: "mkdir", Path: dir, Err: syscall.ENOTDIR} + } + if !os.IsNotExist(err) { + return err + } + if parent == dir { + return err + } + if err := makeStoreDirectory(parent, syncParent); err != nil { + return err + } + if err := os.Mkdir(dir, 0o700); err != nil { + if os.IsExist(err) { + if info, statErr := os.Stat(dir); statErr == nil && info.IsDir() { + return syncParent(parent) + } + } + return err + } + return syncParent(parent) +} + +func syncStoreDirectory(dir string) error { + file, err := os.Open(dir) + if err != nil { + return err + } + err = file.Sync() + closeErr := file.Close() + if err != nil { + return err + } + return closeErr +} diff --git a/store_directory_retry_test.go b/store_directory_retry_test.go new file mode 100644 index 0000000..a149dea --- /dev/null +++ b/store_directory_retry_test.go @@ -0,0 +1,107 @@ +package main + +import ( + "errors" + "os" + "path/filepath" + "syscall" + "testing" +) + +func TestStoreDirectoryRetriesFailedParentBarrier(t *testing.T) { + for _, failure := range []string{"first", "second", "ledger"} { + t.Run(failure, func(t *testing.T) { + root := t.TempDir() + paths := []string{filepath.Join(root, "first"), filepath.Join(root, "first", "second"), filepath.Join(root, "first", "second", "ledger")} + dir := paths[len(paths)-1] + var failedPath string + for _, path := range paths { + if filepath.Base(path) == failure { + failedPath = path + } + } + failedParent := filepath.Dir(failedPath) + failedCalls := 0 + for attempt := range 2 { + lock, _, err := openStoreLock(dir, "retry.lock", func(parent string) error { + if parent == failedParent { + failedCalls++ + return syscall.EIO + } + return syncStoreDirectory(parent) + }) + if lock != nil { + _ = lock.Close() + t.Fatal("retry initialized ownership without completing the failed parent barrier") + } + if !errors.Is(err, syscall.EIO) || failedCalls != attempt+1 { + t.Fatal("retry skipped the failed directory entry's durability barrier") + } + info, statErr := os.Lstat(failedPath) + if statErr != nil || !info.IsDir() || info.Mode().Perm() != 0o700 { + t.Fatal("failed barrier did not preserve its private directory for retry") + } + for _, name := range []string{"retry.lock", "state.json"} { + if _, statErr := os.Lstat(filepath.Join(dir, name)); !os.IsNotExist(statErr) { + t.Fatal("failed retry created an ownership record") + } + } + } + retried := false + lock, created, err := openStoreLock(dir, "retry.lock", func(parent string) error { + if parent == failedParent { + retried = true + } + if _, statErr := os.Lstat(filepath.Join(dir, "retry.lock")); !os.IsNotExist(statErr) { + t.Error("initialization witness preceded retry durability") + } + return syncStoreDirectory(parent) + }) + if err != nil || lock == nil || !created || !retried { + t.Fatal("successful parent barrier could not resume initialization") + } + if err := lock.Close(); err != nil { + t.Fatal("could not close initialized lock") + } + }) + } +} + +func TestStoreDirectoryCompetingMkdirRequiresParentBarrier(t *testing.T) { + root := t.TempDir() + parent := filepath.Join(root, "first") + dir := filepath.Join(parent, "ledger") + competingMkdir, checked := false, false + lock, _, err := openStoreLock(dir, "retry.lock", func(path string) error { + if path == root { + // The target was absent at entry. Materialize it while the first + // ancestor barrier runs, so its later mkdir takes the EEXIST branch. + if err := os.Mkdir(dir, 0o700); err != nil { + return err + } + competingMkdir = true + } + if path == parent { + checked = true + return syscall.EIO + } + return syncStoreDirectory(path) + }) + if lock != nil { + _ = lock.Close() + t.Fatal("competing mkdir bypassed its parent durability barrier") + } + if !competingMkdir || !checked || !errors.Is(err, syscall.EIO) { + t.Fatal("EEXIST did not require a successful parent barrier") + } + if _, err := os.Lstat(filepath.Join(dir, "retry.lock")); !os.IsNotExist(err) { + t.Fatal("competing mkdir permitted ownership initialization after failed sync") + } + lock, created, err := openStoreLock(dir, "retry.lock", syncStoreDirectory) + if err != nil || lock == nil || !created { + t.Fatal("durable retry after competing mkdir failed") + } + if err := lock.Close(); err != nil { + t.Fatal("could not close initialized lock") + } +} diff --git a/store_directory_test.go b/store_directory_test.go new file mode 100644 index 0000000..707977a --- /dev/null +++ b/store_directory_test.go @@ -0,0 +1,99 @@ +package main + +import ( + "errors" + "os" + "path/filepath" + "reflect" + "syscall" + "testing" +) + +func TestStoreDirectorySyncsNewEntriesBeforeInitialization(t *testing.T) { + root := t.TempDir() + parents := []string{root, filepath.Join(root, "first"), filepath.Join(root, "first", "second")} + dir := filepath.Join(root, "first", "second", "ledger") + var synced []string + lock, created, err := openStoreLock(dir, "fixture.lock", func(parent string) error { + synced = append(synced, parent) + if _, err := os.Lstat(filepath.Join(dir, "fixture.lock")); !os.IsNotExist(err) { + t.Error("initialization witness preceded parent durability") + } + return syncStoreDirectory(parent) + }) + if err != nil { + t.Fatal("could not initialize nested private state directory") + } + if err := lock.Close(); err != nil { + t.Fatal("could not release initialized directory") + } + if !created || !reflect.DeepEqual(synced, append([]string{filepath.Dir(root)}, parents...)) { + t.Fatal("created directory entries were not durably ordered before initialization") + } + for _, path := range append(parents[1:], dir) { + info, err := os.Lstat(path) + if err != nil || !info.IsDir() || info.Mode().Perm() != 0o700 { + t.Fatal("new state directory component was not private") + } + } + synced = nil + lock, created, err = openStoreLock(dir, "fixture.lock", func(parent string) error { + synced = append(synced, parent) + return syncStoreDirectory(parent) + }) + if err != nil { + t.Fatal("existing directory could not recover after parent durability") + } + if err := lock.Close(); err != nil || created { + t.Fatal("existing initialization witness was recreated") + } + if !reflect.DeepEqual(synced, []string{filepath.Dir(dir)}) { + t.Fatal("existing directory skipped its parent durability barrier") + } +} + +func TestStoreDirectoryParentSyncFailurePreventsInitialization(t *testing.T) { + for _, failure := range []string{"first", "second", "ledger"} { + t.Run(failure, func(t *testing.T) { + root := t.TempDir() + paths := []string{filepath.Join(root, "first"), filepath.Join(root, "first", "second"), filepath.Join(root, "first", "second", "ledger")} + parents := []string{filepath.Dir(root), root, paths[0], paths[1]} + dir := paths[len(paths)-1] + calls := 0 + failureIndex := -1 + for i, path := range paths { + if filepath.Base(path) == failure { + failureIndex = i + } + } + lock, _, err := openStoreLock(dir, "fixture.lock", func(parent string) error { + index := calls + calls++ + if parent != parents[index] { + t.Error("unexpected parent sync order") + } + if index == failureIndex+1 { + return syscall.EIO + } + return syncStoreDirectory(parent) + }) + if lock != nil { + _ = lock.Close() + t.Fatal("failed parent sync granted an initialization lock") + } + if !errors.Is(err, syscall.EIO) || calls != failureIndex+2 { + t.Fatal("parent durability failure was ignored or initialization continued") + } + for _, name := range []string{"fixture.lock", "state.json"} { + if _, err := os.Lstat(filepath.Join(dir, name)); !os.IsNotExist(err) { + t.Fatal("failed parent sync created an ownership record") + } + } + for _, path := range paths[failureIndex+1:] { + if _, err := os.Lstat(path); !os.IsNotExist(err) { + t.Fatal("directory creation continued beyond a failed durability barrier") + } + } + }) + } +} diff --git a/store_initialization_test.go b/store_initialization_test.go new file mode 100644 index 0000000..ef132b5 --- /dev/null +++ b/store_initialization_test.go @@ -0,0 +1,185 @@ +package main + +import ( + "os" + "path/filepath" + "syscall" + "testing" +) + +type storeInitializationFixture struct { + lockName string + digest string + open func(string) (func(), string, error) + save func(string) error +} + +func storeInitializationFixtures(t *testing.T) map[string]storeInitializationFixture { + t.Helper() + digest := brokerSHA([]byte("durable-broker-fixture")) + c := brokerTestContext(brokerConfiguration{configDigest: digest}) + brokerLedger := &brokerLedger{Version: 1, ConfigDigest: digest, Sessions: map[string]*brokerSession{ + brokerJSONDigest(c.Owner): {Owner: c.Owner, CreateState: "none", Retiring: true, Retired: true, + ProofDigest: brokerSHA([]byte("retired-owner")), Prompts: map[string]*brokerPrompt{}, + Responses: map[string]brokerResponseID{}, Operations: map[string]string{}}, + }} + if !brokerLedgerValid(brokerLedger, digest) { + t.Fatal("invalid retired broker fixture") + } + cfg, exposed := hostedBoundaryLedgerFixture(t) + return map[string]storeInitializationFixture{ + "broker": { + lockName: "broker.lock", digest: brokerJSONDigest(brokerLedger), + open: func(dir string) (func(), string, error) { + store, ledger, err := openBrokerStore(dir, digest) + if err != nil { + return nil, "", err + } + return store.close, brokerJSONDigest(ledger), nil + }, + save: func(dir string) error { return (&brokerStore{dir: dir}).save(brokerLedger) }, + }, + "gateway": { + lockName: "gateway.lock", digest: brokerJSONDigest(exposed), + open: func(dir string) (func(), string, error) { + store, ledger, err := openHostedGatewayStore(dir, cfg) + if err != nil { + return nil, "", err + } + return store.close, brokerJSONDigest(ledger), nil + }, + save: func(dir string) error { return (&hostedGatewayStore{dir: dir}).save(exposed) }, + }, + } +} + +func TestDurableStoreMissingLedgerFailsClosed(t *testing.T) { + for name, fixture := range storeInitializationFixtures(t) { + t.Run(name, func(t *testing.T) { + dir := filepath.Join(t.TempDir(), "ledger") + closeStore, _, err := fixture.open(dir) + if err != nil { + t.Fatal("could not initialize ownership fixture") + } + err = fixture.save(dir) + closeStore() + if err != nil { + t.Fatal("could not persist owned lifetime") + } + lockPath := filepath.Join(dir, fixture.lockName) + lockInfo, err := os.Lstat(lockPath) + if err != nil { + t.Fatal("ownership witness is missing") + } + statePath := filepath.Join(dir, "state.json") + retained := filepath.Join(t.TempDir(), "retained-state.json") + if os.Rename(statePath, retained) != nil { + t.Fatal("could not simulate missing ownership state") + } + for range 2 { + closeStore, _, err = fixture.open(dir) + if err == nil { + closeStore() + t.Fatal("missing ownership ledger was silently reinitialized") + } + if _, err := os.Lstat(statePath); !os.IsNotExist(err) { + t.Fatal("rejected recovery recreated state") + } + after, err := os.Lstat(lockPath) + if err != nil || !os.SameFile(lockInfo, after) { + t.Fatal("rejected recovery replaced its initialization witness") + } + } + if os.Rename(retained, statePath) != nil { + t.Fatal("could not restore exact original ownership") + } + closeStore, restored, err := fixture.open(dir) + if err != nil { + t.Fatal("valid original ownership could not recover") + } + closeStore() + if restored != fixture.digest { + t.Fatal("recovery changed the owned lifetime") + } + }) + } +} + +func TestDurableStoreEmptyDirectoryAndLegacyRecovery(t *testing.T) { + for name, fixture := range storeInitializationFixtures(t) { + t.Run(name, func(t *testing.T) { + dir := filepath.Join(t.TempDir(), "ledger") + if os.Mkdir(dir, 0o700) != nil { + t.Fatal("could not create private pre-existing directory") + } + closeStore, _, err := fixture.open(dir) + if err != nil { + t.Fatal("pre-existing empty directory could not initialize") + } + err = fixture.save(dir) + closeStore() + if err != nil { + t.Fatal("could not persist owned lifetime") + } + for _, missingLock := range []bool{false, true} { + if missingLock && os.Remove(filepath.Join(dir, fixture.lockName)) != nil { + t.Fatal("could not prepare valid ledger without a lock") + } + closeStore, restored, err := fixture.open(dir) + if err != nil { + t.Fatal("valid existing ownership could not recover") + } + closeStore() + if restored != fixture.digest { + t.Fatal("existing ownership was rebound during recovery") + } + } + }) + } +} + +func TestDurableStoreInitializerRace(t *testing.T) { + for name, fixture := range storeInitializationFixtures(t) { + t.Run(name, func(t *testing.T) { + dir := filepath.Join(t.TempDir(), "ledger") + if os.Mkdir(dir, 0o700) != nil { + t.Fatal("could not create private initializer directory") + } + // Pause the original creator after exclusive lock-file creation, + // before flock. A competing opener can acquire flock first, but + // it must not claim the original creator's initialization rights. + witness, err := os.OpenFile(filepath.Join(dir, fixture.lockName), os.O_CREATE|os.O_EXCL|os.O_RDWR|syscall.O_NOFOLLOW, 0o600) + if err != nil { + t.Fatal("could not pause original initialization") + } + defer witness.Close() + closeStore, _, err := fixture.open(dir) + if err == nil { + closeStore() + t.Fatal("competing opener initialized another process's store") + } + if _, err := os.Lstat(filepath.Join(dir, "state.json")); !os.IsNotExist(err) { + t.Fatal("competing opener published a ledger") + } + // Resume the original creator using its same descriptor. The + // rejected contender must release flock and preserve that inode. + if syscall.Flock(int(witness.Fd()), syscall.LOCK_EX|syscall.LOCK_NB) != nil { + t.Fatal("rejected contender retained the initialization lock") + } + if fixture.save(dir) != nil { + t.Fatal("original creator could not publish its ledger") + } + if syscall.Flock(int(witness.Fd()), syscall.LOCK_UN) != nil { + t.Fatal("original creator could not release the store") + } + closeStore, restored, err := fixture.open(dir) + if err != nil { + t.Fatal("competing opener could not read the completed original store") + } + closeStore() + if restored != fixture.digest { + t.Fatal("competing opener replaced original ownership") + } + }) + } +} diff --git a/store_initializer_race_test.go b/store_initializer_race_test.go new file mode 100644 index 0000000..a902c36 --- /dev/null +++ b/store_initializer_race_test.go @@ -0,0 +1,224 @@ +package main + +import ( + "bytes" + "errors" + "os" + "path/filepath" + "syscall" + "testing" + "time" +) + +func TestStoreInitializerKeepsCreatorAcrossFlockContention(t *testing.T) { + for name, fixture := range storeInitializationFixtures(t) { + t.Run(name, func(t *testing.T) { + dir := filepath.Join(t.TempDir(), "ledger") + if err := os.Mkdir(dir, 0o700); err != nil { + t.Fatal("could not create private initializer directory") + } + path := filepath.Join(dir, fixture.lockName) + // Pause creator A after its successful O_EXCL, before the exact + // acquisition helper used by openStoreLock. B opens A's same inode. + creator, err := os.OpenFile(path, os.O_CREATE|os.O_EXCL|os.O_RDWR|syscall.O_NOFOLLOW, 0o600) + if err != nil { + t.Fatal("could not create initialization witness") + } + defer creator.Close() + original, err := creator.Stat() + if err != nil { + t.Fatal("could not inspect creator witness") + } + contender, created, err := openStoreLock(dir, fixture.lockName, syncStoreDirectory) + if err != nil || contender == nil || created { + t.Fatal("contender did not acquire the existing initialization inode") + } + defer contender.Close() + contenderInfo, err := contender.Stat() + if err != nil || !os.SameFile(original, contenderInfo) { + t.Fatal("contender did not lock the creator's exact inode") + } + if _, err := os.Lstat(filepath.Join(dir, "state.json")); !os.IsNotExist(err) { + t.Fatal("contender found unexpected ownership state") + } + started := make(chan struct{}) + acquired := make(chan error, 1) + go func() { + close(started) + acquired <- flockStoreFile(creator, true) + }() + <-started + select { + case err := <-acquired: + // This is the original bug: A abandons O_EXCL authority while + // B is still checking the absent ledger. Neither can initialize. + _ = contender.Close() + _ = creator.Close() + for range 2 { + closeStore, _, openErr := fixture.open(dir) + if openErr == nil { + closeStore() + t.Fatal("a later opener fabricated absent initialization authority") + } + } + if errors.Is(err, syscall.EWOULDBLOCK) { + t.Fatal("exclusive creator lost authority to transient flock contention; later openers remain stranded") + } + t.Fatal("creator completed acquisition while a contender still held the inode") + case <-time.After(75 * time.Millisecond): + // B has no O_EXCL authority. Its caller rejects the absent + // ledger and closes; A must keep its descriptor while waiting. + } + if err := contender.Close(); err != nil { + t.Fatal("could not release the rejected contender") + } + select { + case err := <-acquired: + if err != nil { + t.Fatal("original creator could not resume after contender release") + } + case <-time.After(2 * time.Second): + t.Fatal("original creator remained blocked after contender release") + } + after, err := creator.Stat() + pathInfo, pathErr := os.Lstat(path) + if err != nil || pathErr != nil || !os.SameFile(original, after) || !os.SameFile(original, pathInfo) { + t.Fatal("creator replaced or reopened the initialization witness") + } + if err := fixture.save(dir); err != nil { + t.Fatal("retained creator could not publish original ownership") + } + if err := creator.Close(); err != nil { + t.Fatal("could not close initialized creator") + } + closeStore, restored, err := fixture.open(dir) + if err != nil { + t.Fatal("completed original initialization could not recover") + } + closeStore() + if restored != fixture.digest { + t.Fatal("recovery replaced the original ownership ledger") + } + }) + } +} + +func TestStoreInitializerExistingWriterRemainsNonblocking(t *testing.T) { + for name, fixture := range storeInitializationFixtures(t) { + t.Run(name, func(t *testing.T) { + dir := filepath.Join(t.TempDir(), "ledger") + closeWriter, _, err := fixture.open(dir) + if err != nil { + t.Fatal("could not open original writer") + } + defer closeWriter() + if fixture.save(dir) != nil { + t.Fatal("could not preserve original ownership") + } + finished := make(chan error, 1) + go func() { + other, created, err := openStoreLock(dir, fixture.lockName, syncStoreDirectory) + if other != nil { + _ = other.Close() + } + if other != nil || created { + finished <- errors.New("existing writer was not excluded") + return + } + finished <- err + }() + select { + case err := <-finished: + if !errors.Is(err, syscall.EWOULDBLOCK) { + t.Fatal("existing-inode writer did not fail with nonblocking contention") + } + case <-time.After(2 * time.Second): + t.Fatal("existing-inode writer waited for the active owner") + } + closeWriter() + closeRecovered, digest, err := fixture.open(dir) + if err != nil { + t.Fatal("original ownership did not recover after active writer closed") + } + closeRecovered() + if digest != fixture.digest { + t.Fatal("excluded writer changed original ownership") + } + }) + } +} + +func TestStoreInitializerLegacyRecoveryRemainsNonblocking(t *testing.T) { + for name, fixture := range storeInitializationFixtures(t) { + t.Run(name, func(t *testing.T) { + dir := filepath.Join(t.TempDir(), "ledger") + if os.Mkdir(dir, 0o700) != nil || fixture.save(dir) != nil { + t.Fatal("could not prepare valid legacy ownership without a lock") + } + statePath := filepath.Join(dir, "state.json") + before, err := os.ReadFile(statePath) + if err != nil { + t.Fatal("could not inspect original legacy ownership") + } + path := filepath.Join(dir, fixture.lockName) + // A creates the missing lock for a valid legacy ledger, then B + // recovers that ledger on the same inode before A calls flock. + creator, err := os.OpenFile(path, os.O_CREATE|os.O_EXCL|os.O_RDWR|syscall.O_NOFOLLOW, 0o600) + if err != nil { + t.Fatal("could not create the legacy recovery lock") + } + defer creator.Close() + original, err := creator.Stat() + if err != nil { + t.Fatal("could not inspect the legacy recovery lock") + } + closeWriter, restored, err := fixture.open(dir) + if err != nil || restored != fixture.digest { + if closeWriter != nil { + closeWriter() + } + t.Fatal("contender could not recover the original legacy ledger") + } + defer closeWriter() + contenderInfo, err := os.Lstat(path) + if err != nil || !os.SameFile(original, contenderInfo) { + t.Fatal("contender replaced the creator's lock inode") + } + finished := make(chan error, 1) + go func() { finished <- flockStoreFile(creator, true) }() + select { + case err := <-finished: + if !errors.Is(err, syscall.EWOULDBLOCK) { + t.Fatal("new legacy recovery lock did not fail with nonblocking contention") + } + case <-time.After(2 * time.Second): + // Release B only after proving A waited behind a valid owner. + // Join A so the regression never leaks a blocked syscall. + closeWriter() + select { + case <-finished: + case <-time.After(2 * time.Second): + t.Fatal("legacy recovery acquisition did not finish after owner release") + } + t.Fatal("newly created legacy recovery lock waited behind the retained ledger owner") + } + after, err := os.ReadFile(statePath) + if err != nil || !bytes.Equal(before, after) { + t.Fatal("excluded legacy recovery changed original ownership") + } + pathInfo, err := os.Lstat(path) + if err != nil || !os.SameFile(original, pathInfo) { + t.Fatal("excluded legacy recovery replaced the lock inode") + } + closeWriter() + closeRecovered, restored, err := fixture.open(dir) + if err != nil { + t.Fatal("legacy ownership could not recover after the active owner closed") + } + closeRecovered() + if restored != fixture.digest { + t.Fatal("legacy recovery changed the owned lifetime") + } + }) + } +}