From 8e6454307cf912e483c0de8f0be04ceaa03269cb Mon Sep 17 00:00:00 2001 From: Zishu Yu Date: Sat, 13 Jun 2026 22:04:56 +0800 Subject: [PATCH 01/40] docs: design linux headless server --- ...2026-06-13-linux-headless-server-design.md | 276 ++++++++++++++++++ 1 file changed, 276 insertions(+) create mode 100644 docs/superpowers/specs/2026-06-13-linux-headless-server-design.md diff --git a/docs/superpowers/specs/2026-06-13-linux-headless-server-design.md b/docs/superpowers/specs/2026-06-13-linux-headless-server-design.md new file mode 100644 index 0000000..4f75ebd --- /dev/null +++ b/docs/superpowers/specs/2026-06-13-linux-headless-server-design.md @@ -0,0 +1,276 @@ +# Linux Headless Server Design + +## Goal + +Build a Linux server edition for machines without a desktop environment. The +package is a single extracted directory and exposes one user-facing binary, +`agentserver`. Running `agentserver` in a directory starts that directory as a +foreground slave. Running `agentserver install-driver` configures the current +user's Codex CLI to automatically mount the driver MCP server. + +The Linux path must not change the Windows installer, launcher, or desktop +flows. + +## Scope + +- Linux `amd64` and `arm64` tarballs. +- One main CLI binary named `agentserver`. +- Foreground directory slave execution. +- Driver installation and remote workspace switching. +- Code-side modelserver authentication through device code. +- Agentserver-side device auth for driver and first-time slave joins. +- Local token proxy/refresher only when the user is not already using a + long-lived modelserver API key. +- No systemd service in the first version. +- No desktop browser requirement. + +## Confirmed Auth Facts + +Linux code login uses Hydra's native device flow on `codeapi.cs.ac.cn`, not the +modelserver repository's `/oauth/device/code` wrapper. The wrapper exists in the +source tree but is not exposed by the current production deployment. + +Use the existing public modelserver client id already used by the Windows PKCE +flow: + +```text +5321f7e6-3d79-4ac9-a742-04809dbf9025 +``` + +The code-side device flow uses: + +- authorization endpoint: `https://codeapi.cs.ac.cn/oauth2/device/auth` +- token endpoint: `https://codeapi.cs.ac.cn/oauth2/token` +- scope: `project:inference offline_access` +- grant type: `urn:ietf:params:oauth:grant-type:device_code` +- refresh grant: `refresh_token` + +The returned `access_token` and `refresh_token` are opaque to the client. The +client stores and refreshes them, but never parses token internals. The upstream +modelserver proxy introspects bearer tokens and reads extension claims such as +`project_id` and `user_id`; that interpretation stays server-side. + +## Modelserver Access Rule + +Every Linux `agentserver` command starts by ensuring modelserver access. The +check has two modes, in this order: + +1. Prefer an existing long-lived API key. If `OPENAI_API_KEY` is set in the + environment, repair `~/.codex/config.toml` to direct modelserver mode if + needed, do not start the local proxy/refresher, and do not require code + device login: + + ```toml + [model_providers.modelserver] + name = "modelserver" + base_url = "https://code.ai.cs.ac.cn/v1" + env_key = "OPENAI_API_KEY" + wire_api = "responses" + ``` + +2. Otherwise, ensure device-code credentials and a single user-level local + proxy/refresher. The Codex modelserver provider is repaired to: + + ```toml + [model_providers.modelserver] + name = "modelserver" + base_url = "http://127.0.0.1:53452/v1" + env_key = "AGENTSERVER_CODEX_LOCAL_API_KEY" + wire_api = "responses" + ``` + +In proxy mode, the local proxy injects the latest access token from the secrets +store for every request to `https://code.ai.cs.ac.cn/v1`. The refresher uses the +stored refresh token and marks reauth-required on `invalid_grant` or missing +refresh token. The next user-facing command reruns code device login. + +This rule applies to `agentserver`, `agentserver install-driver`, +`agentserver switch-workspace`, `agentserver serve-driver-mcp`, and future +diagnostic/login subcommands. + +## Commands + +### `agentserver` + +The default command runs the current directory as a slave in the foreground. + +Startup sequence: + +1. Run the modelserver access rule. +2. Resolve the current directory to a canonical path. +3. Look up the canonical path in the slave registry. +4. If absent, prompt for a slave name. The default is + `-`. +5. Write or repair the slave config for that directory. +6. Start `slave-agent` as a foreground child. +7. Capture auth URLs emitted by `slave-agent`. +8. If an `agent.cs.ac.cn` device URL appears, print the URL, user code when + available, and an ASCII QR code. +9. Exit when the child exits. Ctrl-C stops the child. + +The same directory always reuses the same slave identity and config unless the +registry is manually removed. + +### `agentserver install-driver` + +Configures the current user as a driver for ordinary Codex CLI sessions. + +Startup sequence: + +1. Run the modelserver access rule. +2. Ensure agentserver workspace credentials. If missing or invalid, run + agentserver device auth and print the URL/code/QR. +3. Register or repair the driver config. The display name is + `-星池指挥官`. +4. Register the Codex MCP server globally in `~/.codex/config.toml`: + + ```toml + [mcp_servers.driver] + command = "/path/to/agentserver" + args = ["serve-driver-mcp"] + ``` + +5. Health-check the resulting setup and print the current remote workspace. + +This command is idempotent. Re-running it repairs missing config and leaves the +selected workspace unchanged. + +### `agentserver switch-workspace` + +Reruns agentserver-side device auth so the user can select a different remote +workspace. Then it rewrites the driver config and keeps the Codex MCP entry +point unchanged. + +The driver has no local directory concept. Its local workdir follows the Codex +CLI process that starts the MCP server. The only workspace switch here is the +remote agentserver workspace. + +### `agentserver serve-driver-mcp` + +Internal subcommand used by Codex MCP. It runs the modelserver access rule, +loads the current driver config, and then starts or execs +`driver-agent serve-mcp`. + +If long-lived API key mode is active, it does not start the local +proxy/refresher. Otherwise it ensures the user-level singleton proxy/refresher +is running before serving MCP. + +## Auth UX + +Code-side auth prints: + +- verification URL from Hydra, usually + `https://codeapi.cs.ac.cn/oauth2/device/verify?...` +- user code +- ASCII QR code for the verification URL + +Agentserver-side auth for drivers and slaves prints: + +- `agent.cs.ac.cn` verification URL +- user code when present +- ASCII QR code + +The CLI never attempts to open a browser on Linux. + +## State And Secrets + +All Linux user state lives under the existing `~/.agentserver-app` root unless +the existing path package is configured otherwise. + +- `state.json`: driver workspace metadata, install status, and reauth flags. +- `machine.json`: existing machine identity. +- `slaves.json`: existing slave registry extended with + `canonical_path -> slave_id` lookup. +- `secrets.json`: file fallback for secrets when the platform keyring is not + available; permissions must be `0600`. +- `token-refresher.lock`: user-level singleton lock for the proxy/refresher. +- `bin/` or the extracted package directory: `agentserver`, `driver-agent`, + `slave-agent`, and the optional pinned Codex runtime. + +Linux secrets use the existing `internal/secrets` abstraction. Keyring is used +when available. If keyring setup fails or is unavailable on a headless server, +the CLI falls back to the `0600` file store. + +## Packaging + +Publish two tarballs: + +- `linux-amd64` +- `linux-arm64` + +Each tarball contains: + +- `agentserver` +- pinned `driver-agent` +- pinned `slave-agent` +- metadata with versions and SHA256 checksums + +The packaging script pulls `driver-agent` and `slave-agent` from the pinned Loom +release for the target architecture and verifies SHA256. Missing assets or +checksum mismatches fail the package build clearly. + +Codex runtime handling: + +1. Prefer `codex` already on `PATH`. +2. If missing, download or install the pinned runtime into the app-managed bin + root used by existing paths. + +## Error Handling + +- If `OPENAI_API_KEY` is set but Codex config is not direct modelserver, repair + Codex config to direct modelserver and do not start proxy/refresher. +- If no long-lived key exists and refresh is impossible, rerun code device + login. +- If proxy port `127.0.0.1:53452` is occupied by this user's healthy proxy, + reuse it. +- If the port is occupied by an unhealthy or unrelated process, fail with a + clear error that names the port and suggests stopping the process. +- If a slave directory cannot be canonicalized, fail before writing registry + state. +- If a child `slave-agent` or `driver-agent` is missing, fail with a message + that names the expected path and package architecture. +- If device-code auth expires, print a concise expiry message and let the user + rerun the same command. + +## Testing + +Unit tests: + +- modelserver access mode chooses long-lived API key before proxy/refresher. +- modelserver device config uses Hydra native endpoints and the existing public + client id. +- Codex config merge preserves unrelated settings while updating modelserver + and driver MCP entries. +- slave registry maps canonical directory path to stable slave id. +- secrets fallback creates `0600` files. +- proxy/refresher lock allows one process and rejects duplicates. + +Integration tests with fakes: + +- fake codeapi device flow returns challenge, pending, slow_down, approved, and + invalid_grant refresh cases. +- fake agentserver device flow returns driver and slave workspace credentials. +- fake `slave-agent` emits an auth URL; `agentserver` captures and prints it. +- fake `driver-agent serve-mcp` is launched through `serve-driver-mcp`. + +Build and packaging tests: + +- `GOOS=linux GOARCH=amd64 go test ./...` +- `GOOS=linux GOARCH=arm64 go test ./...` +- package dry-run validates asset selection, SHA256 checking, and archive + layout. + +Manual online verification: + +- create a code device challenge against production and confirm pending token + polling semantics. +- do not automate real user login in CI. + +## Non-Goals + +- No Linux desktop launcher. +- No systemd service. +- No root install path. +- No changes to Windows installer UX. +- No dependency on modelserver's currently unexposed `/oauth/device/code` + wrapper. From 862cc7af8bc3fefec0afb9ba4c9e83582e920776 Mon Sep 17 00:00:00 2001 From: Zishu Yu Date: Sat, 13 Jun 2026 22:47:46 +0800 Subject: [PATCH 02/40] feat(headless): add Linux device auth configs --- cmd/launcher/main.go | 17 ++------------- internal/agentserver/oauth.go | 20 ++++++++++++++++++ internal/agentserver/oauth_test.go | 29 ++++++++++++++++++++++++++ internal/modelserver/device.go | 21 +++++++++++++++++++ internal/modelserver/device_test.go | 32 +++++++++++++++++++++++++++++ 5 files changed, 104 insertions(+), 15 deletions(-) create mode 100644 internal/agentserver/oauth.go create mode 100644 internal/agentserver/oauth_test.go create mode 100644 internal/modelserver/device.go create mode 100644 internal/modelserver/device_test.go diff --git a/cmd/launcher/main.go b/cmd/launcher/main.go index d382302..e327896 100644 --- a/cmd/launcher/main.go +++ b/cmd/launcher/main.go @@ -563,7 +563,7 @@ func newCompletedConsoleOrchestrator(in completedOrchestratorInput) ui.Orchestra asBaseURL := completedAgentserverBaseURL(in.State) asOAuth := in.ASOAuth if asOAuth.Endpoint == "" { - asOAuth = defaultAgentserverOAuthConfig(asBaseURL) + asOAuth = agentserver.OAuthConfig(asBaseURL) } loomDriverPath := "" if in.InstallDir != "" { @@ -678,7 +678,7 @@ func serveOnboarding(p paths.Paths, store *state.Store) error { // to Hydra. The CLI client `agentserver-agent-cli` is pre-registered // by the Helm chart with grant=device_code, public (no secret), // scopes=openid profile agent:register. - asOAuth := defaultAgentserverOAuthConfig("https://agent.cs.ac.cn") + asOAuth := agentserver.OAuthConfig("https://agent.cs.ac.cn") installDir, err := os.Executable() if err != nil { @@ -749,19 +749,6 @@ func serveOnboarding(p paths.Paths, store *state.Store) error { return err } -func defaultAgentserverOAuthConfig(endpoint string) oauth.Config { - if strings.TrimSpace(endpoint) == "" { - endpoint = "https://agent.cs.ac.cn" - } - return oauth.Config{ - Endpoint: strings.TrimRight(strings.TrimSpace(endpoint), "/"), - AuthPath: "/api/oauth2/device/auth", - TokenPath: "/api/oauth2/token", - ClientID: "agentserver-agent-cli", - Scope: "openid profile agent:register", - } -} - func launchCompletedInstall(ctx context.Context, codeExe string, p paths.Paths, sec secrets.Store, tokenRefresherExe string, embeddedVSIXPath string) error { if err := launchprep.PrepareVSCode(ctx, launchprep.Input{ CodeExe: codeExe, diff --git a/internal/agentserver/oauth.go b/internal/agentserver/oauth.go new file mode 100644 index 0000000..e2bf22e --- /dev/null +++ b/internal/agentserver/oauth.go @@ -0,0 +1,20 @@ +package agentserver + +import ( + "strings" + + "github.com/agentserver/agentserver-pkg/internal/oauth" +) + +func OAuthConfig(endpoint string) oauth.Config { + if strings.TrimSpace(endpoint) == "" { + endpoint = "https://agent.cs.ac.cn" + } + return oauth.Config{ + Endpoint: strings.TrimRight(strings.TrimSpace(endpoint), "/"), + AuthPath: "/api/oauth2/device/auth", + TokenPath: "/api/oauth2/token", + ClientID: "agentserver-agent-cli", + Scope: "openid profile agent:register", + } +} diff --git a/internal/agentserver/oauth_test.go b/internal/agentserver/oauth_test.go new file mode 100644 index 0000000..dd8fb8e --- /dev/null +++ b/internal/agentserver/oauth_test.go @@ -0,0 +1,29 @@ +package agentserver + +import "testing" + +func TestOAuthConfigUsesPublicAgentCLIClient(t *testing.T) { + got := OAuthConfig("") + if got.Endpoint != "https://agent.cs.ac.cn" { + t.Fatalf("Endpoint=%q", got.Endpoint) + } + if got.AuthPath != "/api/oauth2/device/auth" { + t.Fatalf("AuthPath=%q", got.AuthPath) + } + if got.TokenPath != "/api/oauth2/token" { + t.Fatalf("TokenPath=%q", got.TokenPath) + } + if got.ClientID != "agentserver-agent-cli" { + t.Fatalf("ClientID=%q", got.ClientID) + } + if got.Scope != "openid profile agent:register" { + t.Fatalf("Scope=%q", got.Scope) + } +} + +func TestOAuthConfigTrimsEndpoint(t *testing.T) { + got := OAuthConfig("https://agent.test/") + if got.AuthURL() != "https://agent.test/api/oauth2/device/auth" { + t.Fatalf("AuthURL=%q", got.AuthURL()) + } +} diff --git a/internal/modelserver/device.go b/internal/modelserver/device.go new file mode 100644 index 0000000..37ad2b6 --- /dev/null +++ b/internal/modelserver/device.go @@ -0,0 +1,21 @@ +package modelserver + +import ( + "strings" + + "github.com/agentserver/agentserver-pkg/internal/oauth" +) + +func DeviceConfig(endpoint string) oauth.Config { + cfg := OAuthConfig() + if strings.TrimSpace(endpoint) == "" { + endpoint = cfg.Endpoint + } + return oauth.Config{ + Endpoint: strings.TrimRight(strings.TrimSpace(endpoint), "/"), + AuthPath: "/oauth2/device/auth", + TokenPath: cfg.TokenPath, + ClientID: cfg.ClientID, + Scope: cfg.Scope, + } +} diff --git a/internal/modelserver/device_test.go b/internal/modelserver/device_test.go new file mode 100644 index 0000000..9eed5a7 --- /dev/null +++ b/internal/modelserver/device_test.go @@ -0,0 +1,32 @@ +package modelserver + +import "testing" + +func TestDeviceConfigUsesHydraNativeDeviceFlow(t *testing.T) { + got := DeviceConfig("") + if got.Endpoint != "https://codeapi.cs.ac.cn" { + t.Fatalf("Endpoint=%q", got.Endpoint) + } + if got.AuthPath != "/oauth2/device/auth" { + t.Fatalf("AuthPath=%q", got.AuthPath) + } + if got.TokenPath != "/oauth2/token" { + t.Fatalf("TokenPath=%q", got.TokenPath) + } + if got.ClientID != "5321f7e6-3d79-4ac9-a742-04809dbf9025" { + t.Fatalf("ClientID=%q", got.ClientID) + } + if got.Scope != "project:inference offline_access" { + t.Fatalf("Scope=%q", got.Scope) + } +} + +func TestDeviceConfigAllowsEndpointOverride(t *testing.T) { + got := DeviceConfig("https://codeapi.test/") + if got.Endpoint != "https://codeapi.test" { + t.Fatalf("Endpoint=%q", got.Endpoint) + } + if got.AuthURL() != "https://codeapi.test/oauth2/device/auth" { + t.Fatalf("AuthURL=%q", got.AuthURL()) + } +} From 9f03f1d6df9e3e72051cc913f013f62ee01b7149 Mon Sep 17 00:00:00 2001 From: Zishu Yu Date: Sat, 13 Jun 2026 22:54:26 +0800 Subject: [PATCH 03/40] feat(headless): print device auth QR codes --- go.mod | 3 ++ go.sum | 4 ++ internal/terminalauth/print.go | 57 ++++++++++++++++++++++++++++ internal/terminalauth/print_test.go | 59 +++++++++++++++++++++++++++++ 4 files changed, 123 insertions(+) create mode 100644 internal/terminalauth/print.go create mode 100644 internal/terminalauth/print_test.go diff --git a/go.mod b/go.mod index a09a693..36800b8 100644 --- a/go.mod +++ b/go.mod @@ -5,6 +5,7 @@ go 1.26.3 require ( github.com/BurntSushi/toml v1.4.0 github.com/google/uuid v1.6.0 + github.com/mdp/qrterminal/v3 v3.2.1 golang.org/x/sys v0.45.0 ) @@ -19,7 +20,9 @@ require ( github.com/mattn/go-isatty v0.0.20 // indirect github.com/ncruces/go-strftime v1.0.0 // indirect github.com/remyoudompheng/bigfft v0.0.0-20230129092748-24d4a6f8daec // indirect + golang.org/x/term v0.41.0 // indirect modernc.org/libc v1.72.3 // indirect modernc.org/mathutil v1.7.1 // indirect modernc.org/memory v1.11.0 // indirect + rsc.io/qr v0.2.0 // indirect ) diff --git a/go.sum b/go.sum index 34122da..5f318e5 100644 --- a/go.sum +++ b/go.sum @@ -10,6 +10,8 @@ github.com/hashicorp/golang-lru/v2 v2.0.7 h1:a+bsQ5rvGLjzHuww6tVxozPZFVghXaHOwFs github.com/hashicorp/golang-lru/v2 v2.0.7/go.mod h1:QeFd9opnmA6QUJc5vARoKUSoFhyfM2/ZepoAG6RGpeM= github.com/mattn/go-isatty v0.0.20 h1:xfD0iDuEKnDkl03q4limB+vH+GxLEtL/jb4xVJSWWEY= github.com/mattn/go-isatty v0.0.20/go.mod h1:W+V8PltTTMOvKvAeJH7IuucS94S2C6jfK/D7dTCTo3Y= +github.com/mdp/qrterminal/v3 v3.2.1 h1:6+yQjiiOsSuXT5n9/m60E54vdgFsw0zhADHhHLrFet4= +github.com/mdp/qrterminal/v3 v3.2.1/go.mod h1:jOTmXvnBsMy5xqLniO0R++Jmjs2sTm9dFSuQ5kpz/SU= github.com/ncruces/go-strftime v1.0.0 h1:HMFp8mLCTPp341M/ZnA4qaf7ZlsbTc+miZjCLOFAw7w= github.com/ncruces/go-strftime v1.0.0/go.mod h1:Fwc5htZGVVkseilnfgOVb9mKy6w1naJmn9CehxcKcls= github.com/remyoudompheng/bigfft v0.0.0-20230129092748-24d4a6f8daec h1:W09IVJc94icq4NjY3clb7Lk8O1qJ8BdBEF8z0ibU0rE= @@ -59,3 +61,5 @@ modernc.org/strutil v1.2.1 h1:UneZBkQA+DX2Rp35KcM69cSsNES9ly8mQWD71HKlOA0= modernc.org/strutil v1.2.1/go.mod h1:EHkiggD70koQxjVdSBM3JKM7k6L0FbGE5eymy9i3B9A= modernc.org/token v1.1.0 h1:Xl7Ap9dKaEs5kLoOQeQmPWevfnk/DM5qcLcYlA8ys6Y= modernc.org/token v1.1.0/go.mod h1:UGzOrNV1mAFSEB63lOFHIpNRUVMvYTc6yu1SMY/XTDM= +rsc.io/qr v0.2.0 h1:6vBLea5/NRMVTz8V66gipeLycZMl/+UlFmk8DvqQ6WY= +rsc.io/qr v0.2.0/go.mod h1:IF+uZjkb9fqyeF/4tlBoynqmQxUoPfWEKh921coOuXs= diff --git a/internal/terminalauth/print.go b/internal/terminalauth/print.go new file mode 100644 index 0000000..578dfd3 --- /dev/null +++ b/internal/terminalauth/print.go @@ -0,0 +1,57 @@ +package terminalauth + +import ( + "fmt" + "io" + "strings" + + "github.com/agentserver/agentserver-pkg/internal/oauth" + "github.com/mdp/qrterminal/v3" +) + +type QRWriter func(w interface{ Write([]byte) (int, error) }, url string) + +func ChallengeURL(ch oauth.DeviceCodeChallenge) string { + if strings.TrimSpace(ch.VerificationURIComplete) != "" { + return strings.TrimSpace(ch.VerificationURIComplete) + } + return strings.TrimSpace(ch.VerificationURI) +} + +func PrintChallenge(w io.Writer, title string, ch oauth.DeviceCodeChallenge, qr QRWriter) { + if w == nil { + return + } + if strings.TrimSpace(title) != "" { + fmt.Fprintf(w, "\n%s\n", strings.TrimSpace(title)) + } + url := ChallengeURL(ch) + if url == "" { + return + } + fmt.Fprintf(w, "URL: %s\n", url) + if strings.TrimSpace(ch.UserCode) != "" { + fmt.Fprintf(w, "Code: %s\n", strings.TrimSpace(ch.UserCode)) + } + if qr == nil { + qr = DefaultQR + } + qr(w, url) +} + +func PrintURL(w io.Writer, title, rawURL, userCode string, qr QRWriter) { + PrintChallenge(w, title, oauth.DeviceCodeChallenge{ + UserCode: strings.TrimSpace(userCode), + VerificationURIComplete: strings.TrimSpace(rawURL), + }, qr) +} + +func DefaultQR(w interface{ Write([]byte) (int, error) }, url string) { + qrterminal.GenerateWithConfig(url, qrterminal.Config{ + Level: qrterminal.L, + Writer: w, + BlackChar: qrterminal.BLACK_BLACK, + WhiteChar: qrterminal.WHITE_WHITE, + QuietZone: 1, + }) +} diff --git a/internal/terminalauth/print_test.go b/internal/terminalauth/print_test.go new file mode 100644 index 0000000..93070b0 --- /dev/null +++ b/internal/terminalauth/print_test.go @@ -0,0 +1,59 @@ +package terminalauth + +import ( + "bytes" + "strings" + "testing" + + "github.com/agentserver/agentserver-pkg/internal/oauth" +) + +func TestChallengeURLPrefersCompleteVerificationURI(t *testing.T) { + ch := oauth.DeviceCodeChallenge{ + VerificationURI: "https://example.test/verify", + VerificationURIComplete: "https://example.test/verify?user_code=ABCD", + } + if got := ChallengeURL(ch); got != ch.VerificationURIComplete { + t.Fatalf("ChallengeURL=%q", got) + } +} + +func TestPrintChallengeIncludesTitleURLCodeAndQR(t *testing.T) { + var buf bytes.Buffer + qrCalled := false + ch := oauth.DeviceCodeChallenge{ + UserCode: "ABCD-EFGH", + VerificationURIComplete: "https://codeapi.cs.ac.cn/oauth2/device/verify?user_code=ABCD-EFGH", + } + + PrintChallenge(&buf, "Code 登录", ch, func(w interface{ Write([]byte) (int, error) }, url string) { + qrCalled = true + if url != ch.VerificationURIComplete { + t.Fatalf("QR url=%q", url) + } + _, _ = w.Write([]byte("[qr]\n")) + }) + + out := buf.String() + for _, want := range []string{ + "Code 登录", + "URL: https://codeapi.cs.ac.cn/oauth2/device/verify?user_code=ABCD-EFGH", + "Code: ABCD-EFGH", + "[qr]", + } { + if !strings.Contains(out, want) { + t.Fatalf("output missing %q:\n%s", want, out) + } + } + if !qrCalled { + t.Fatal("QR writer was not called") + } +} + +func TestPrintChallengeSkipsEmptyURL(t *testing.T) { + var buf bytes.Buffer + PrintChallenge(&buf, "Agentserver 登录", oauth.DeviceCodeChallenge{}, nil) + if strings.Contains(buf.String(), "URL:") { + t.Fatalf("unexpected URL output:\n%s", buf.String()) + } +} From 9bb8cde29213b0a4f757bc8ed1d3fe9f1aa97d7c Mon Sep 17 00:00:00 2001 From: Zishu Yu Date: Sat, 13 Jun 2026 22:58:57 +0800 Subject: [PATCH 04/40] fix(headless): harden terminal QR output --- internal/terminalauth/print.go | 26 +++++++++++-- internal/terminalauth/print_test.go | 58 ++++++++++++++++++++++++++++- 2 files changed, 79 insertions(+), 5 deletions(-) diff --git a/internal/terminalauth/print.go b/internal/terminalauth/print.go index 578dfd3..faaaec0 100644 --- a/internal/terminalauth/print.go +++ b/internal/terminalauth/print.go @@ -3,6 +3,7 @@ package terminalauth import ( "fmt" "io" + "reflect" "strings" "github.com/agentserver/agentserver-pkg/internal/oauth" @@ -22,13 +23,13 @@ func PrintChallenge(w io.Writer, title string, ch oauth.DeviceCodeChallenge, qr if w == nil { return } - if strings.TrimSpace(title) != "" { - fmt.Fprintf(w, "\n%s\n", strings.TrimSpace(title)) - } url := ChallengeURL(ch) if url == "" { return } + if strings.TrimSpace(title) != "" { + fmt.Fprintf(w, "\n%s\n", strings.TrimSpace(title)) + } fmt.Fprintf(w, "URL: %s\n", url) if strings.TrimSpace(ch.UserCode) != "" { fmt.Fprintf(w, "Code: %s\n", strings.TrimSpace(ch.UserCode)) @@ -47,6 +48,12 @@ func PrintURL(w io.Writer, title, rawURL, userCode string, qr QRWriter) { } func DefaultQR(w interface{ Write([]byte) (int, error) }, url string) { + if isNilWriter(w) || strings.TrimSpace(url) == "" { + return + } + defer func() { + _ = recover() + }() qrterminal.GenerateWithConfig(url, qrterminal.Config{ Level: qrterminal.L, Writer: w, @@ -55,3 +62,16 @@ func DefaultQR(w interface{ Write([]byte) (int, error) }, url string) { QuietZone: 1, }) } + +func isNilWriter(w interface{ Write([]byte) (int, error) }) bool { + if w == nil { + return true + } + v := reflect.ValueOf(w) + switch v.Kind() { + case reflect.Chan, reflect.Func, reflect.Interface, reflect.Map, reflect.Pointer, reflect.Slice: + return v.IsNil() + default: + return false + } +} diff --git a/internal/terminalauth/print_test.go b/internal/terminalauth/print_test.go index 93070b0..ca48c82 100644 --- a/internal/terminalauth/print_test.go +++ b/internal/terminalauth/print_test.go @@ -18,6 +18,15 @@ func TestChallengeURLPrefersCompleteVerificationURI(t *testing.T) { } } +func TestChallengeURLFallsBackToVerificationURI(t *testing.T) { + ch := oauth.DeviceCodeChallenge{ + VerificationURI: " https://example.test/verify ", + } + if got := ChallengeURL(ch); got != "https://example.test/verify" { + t.Fatalf("ChallengeURL=%q", got) + } +} + func TestPrintChallengeIncludesTitleURLCodeAndQR(t *testing.T) { var buf bytes.Buffer qrCalled := false @@ -53,7 +62,52 @@ func TestPrintChallengeIncludesTitleURLCodeAndQR(t *testing.T) { func TestPrintChallengeSkipsEmptyURL(t *testing.T) { var buf bytes.Buffer PrintChallenge(&buf, "Agentserver 登录", oauth.DeviceCodeChallenge{}, nil) - if strings.Contains(buf.String(), "URL:") { - t.Fatalf("unexpected URL output:\n%s", buf.String()) + if buf.String() != "" { + t.Fatalf("unexpected output:\n%s", buf.String()) } } + +func TestPrintURLTrimsURLAndCode(t *testing.T) { + var buf bytes.Buffer + PrintURL(&buf, "Agentserver 登录", " https://example.test/verify?user_code=ABCD ", " ABCD ", func(w interface{ Write([]byte) (int, error) }, url string) { + if url != "https://example.test/verify?user_code=ABCD" { + t.Fatalf("QR url=%q", url) + } + }) + + out := buf.String() + for _, want := range []string{ + "URL: https://example.test/verify?user_code=ABCD", + "Code: ABCD", + } { + if !strings.Contains(out, want) { + t.Fatalf("output missing %q:\n%s", want, out) + } + } + if strings.Contains(out, " ABCD ") { + t.Fatalf("output contains untrimmed code:\n%s", out) + } +} + +func TestDefaultQRSkipsNilWriter(t *testing.T) { + assertNotPanics(t, func() { + DefaultQR(nil, "https://example.test") + }) +} + +func TestDefaultQRSkipsUnencodableInput(t *testing.T) { + var buf bytes.Buffer + assertNotPanics(t, func() { + DefaultQR(&buf, strings.Repeat("x", 10000)) + }) +} + +func assertNotPanics(t *testing.T, fn func()) { + t.Helper() + defer func() { + if r := recover(); r != nil { + t.Fatalf("unexpected panic: %v", r) + } + }() + fn() +} From 6fa9765859c26cb0c730892963ed778096cb852a Mon Sep 17 00:00:00 2001 From: Zishu Yu Date: Sat, 13 Jun 2026 23:06:16 +0800 Subject: [PATCH 05/40] feat(headless): share model proxy daemon --- cmd/token-refresher/main.go | 59 ++------- internal/modelaccess/daemon.go | 165 ++++++++++++++++++++++++ internal/modelaccess/daemon_test.go | 192 ++++++++++++++++++++++++++++ 3 files changed, 367 insertions(+), 49 deletions(-) create mode 100644 internal/modelaccess/daemon.go create mode 100644 internal/modelaccess/daemon_test.go diff --git a/cmd/token-refresher/main.go b/cmd/token-refresher/main.go index 6e7d822..2a87b1c 100644 --- a/cmd/token-refresher/main.go +++ b/cmd/token-refresher/main.go @@ -7,7 +7,7 @@ import ( "log" "path/filepath" - "github.com/agentserver/agentserver-pkg/internal/modelproxy" + "github.com/agentserver/agentserver-pkg/internal/modelaccess" "github.com/agentserver/agentserver-pkg/internal/modelserver" "github.com/agentserver/agentserver-pkg/internal/oauth" "github.com/agentserver/agentserver-pkg/internal/paths" @@ -55,52 +55,13 @@ func runWithDeps(ctx context.Context, deps runDeps) error { if deps.Logf == nil { deps.Logf = log.Printf } - if deps.LockPath != "" { - lock, err := tokenrefresh.AcquireDaemonLock(deps.LockPath) - if err != nil { - return err - } - defer lock.Close() - } - - proxyErr := make(chan error, 1) - go func() { - proxyErr <- modelproxy.ListenAndServe(ctx, modelproxy.ServerOptions{ - Addr: deps.ProxyAddr, - Secrets: deps.Secrets, - UpstreamBaseURL: deps.ProxyUpstreamBaseURL, - }) - }() - - refreshErr := make(chan error, 1) - go func() { - refreshErr <- tokenrefresh.Run(ctx, tokenrefresh.Options{ - Secrets: deps.Secrets, - OAuth: deps.OAuth, - Refresh: deps.Refresh, - Logf: deps.Logf, - }) - }() - - for { - select { - case err := <-proxyErr: - if ctx.Err() != nil { - return nil - } - return err - case err := <-refreshErr: - if err == nil || errors.Is(err, context.Canceled) { - return nil - } - if errors.Is(err, tokenrefresh.ErrNoRefreshToken) { - deps.Logf("token refresh disabled: %v", err) - refreshErr = nil - continue - } - return err - case <-ctx.Done(): - return nil - } - } + return modelaccess.RunDaemon(ctx, modelaccess.DaemonOptions{ + Secrets: deps.Secrets, + OAuth: deps.OAuth, + ProxyAddr: deps.ProxyAddr, + ProxyUpstreamBaseURL: deps.ProxyUpstreamBaseURL, + LockPath: deps.LockPath, + Refresh: deps.Refresh, + Logf: deps.Logf, + }) } diff --git a/internal/modelaccess/daemon.go b/internal/modelaccess/daemon.go new file mode 100644 index 0000000..e5ef93a --- /dev/null +++ b/internal/modelaccess/daemon.go @@ -0,0 +1,165 @@ +package modelaccess + +import ( + "context" + "errors" + "net/http" + "os/exec" + "path/filepath" + "strings" + "time" + + "github.com/agentserver/agentserver-pkg/internal/modelproxy" + "github.com/agentserver/agentserver-pkg/internal/modelserver" + "github.com/agentserver/agentserver-pkg/internal/oauth" + "github.com/agentserver/agentserver-pkg/internal/process" + "github.com/agentserver/agentserver-pkg/internal/secrets" + "github.com/agentserver/agentserver-pkg/internal/tokenrefresh" +) + +var ErrProxyUnavailable = errors.New("modelaccess: local model proxy unavailable") + +type DaemonOptions struct { + Secrets secrets.Store + OAuth oauth.AuthCodeConfig + ProxyAddr string + ProxyUpstreamBaseURL string + LockPath string + Refresh func(context.Context, oauth.AuthCodeConfig, string) (oauth.Token, error) + Logf func(string, ...any) +} + +func RunDaemon(ctx context.Context, opts DaemonOptions) error { + if ctx == nil { + ctx = context.Background() + } + if opts.Secrets == nil { + return tokenrefresh.ErrNoSecrets + } + if opts.OAuth.ClientID == "" { + opts.OAuth = modelserver.OAuthConfig() + } + if opts.LockPath != "" { + lock, err := tokenrefresh.AcquireDaemonLock(opts.LockPath) + if err != nil { + return err + } + defer lock.Close() + } + + proxyErr := make(chan error, 1) + go func() { + proxyErr <- modelproxy.ListenAndServe(ctx, modelproxy.ServerOptions{ + Addr: opts.ProxyAddr, + Secrets: opts.Secrets, + UpstreamBaseURL: opts.ProxyUpstreamBaseURL, + }) + }() + + refreshErr := make(chan error, 1) + go func() { + refreshErr <- tokenrefresh.Run(ctx, tokenrefresh.Options{ + Secrets: opts.Secrets, + OAuth: opts.OAuth, + Refresh: opts.Refresh, + Logf: opts.Logf, + }) + }() + + for { + select { + case err := <-proxyErr: + if ctx.Err() != nil { + return nil + } + return err + case err := <-refreshErr: + if err == nil || errors.Is(err, context.Canceled) { + return nil + } + if errors.Is(err, tokenrefresh.ErrNoRefreshToken) { + if opts.Logf != nil { + opts.Logf("token refresh disabled: %v", err) + } + refreshErr = nil + continue + } + return err + case <-ctx.Done(): + return nil + } + } +} + +type EnsureDaemonOptions struct { + ExePath string + ProxyBaseURL string + HealthCheck func(context.Context, string) bool + StartProcess func(*exec.Cmd) error +} + +func EnsureDaemon(ctx context.Context, opts EnsureDaemonOptions) error { + if ctx == nil { + ctx = context.Background() + } + baseURL := opts.ProxyBaseURL + if baseURL == "" { + baseURL = "http://" + modelproxy.DefaultListenAddr + } + healthCheck := opts.HealthCheck + if healthCheck == nil { + healthCheck = ProxyHealthy + } + if healthCheck(ctx, baseURL) { + return nil + } + cmd := exec.CommandContext(ctx, opts.ExePath, "model-proxy-daemon") + cmd.Stdin = nil + cmd.Stdout = nil + cmd.Stderr = nil + process.HideWindow(cmd) + startProcess := opts.StartProcess + if startProcess == nil { + startProcess = func(cmd *exec.Cmd) error { return cmd.Start() } + } + if err := startProcess(cmd); err != nil { + if strings.Contains(strings.ToLower(err.Error()), "address already in use") { + return errors.Join(ErrProxyUnavailable, err) + } + return err + } + + deadline := time.Now().Add(3 * time.Second) + for time.Now().Before(deadline) { + if healthCheck(ctx, baseURL) { + return nil + } + select { + case <-ctx.Done(): + return ctx.Err() + case <-time.After(50 * time.Millisecond): + } + } + return ErrProxyUnavailable +} + +func ProxyHealthy(ctx context.Context, baseURL string) bool { + if ctx == nil { + ctx = context.Background() + } + baseURL = strings.TrimRight(baseURL, "/") + req, err := http.NewRequestWithContext(ctx, http.MethodGet, baseURL+modelproxy.HealthPath, nil) + if err != nil { + return false + } + resp, err := http.DefaultClient.Do(req) + if err != nil { + return false + } + defer resp.Body.Close() + return resp.StatusCode == http.StatusNoContent +} + +func DefaultLockPath(installRoot string) string { + return filepath.Join(installRoot, "token-refresher.lock") +} diff --git a/internal/modelaccess/daemon_test.go b/internal/modelaccess/daemon_test.go new file mode 100644 index 0000000..0e85f8e --- /dev/null +++ b/internal/modelaccess/daemon_test.go @@ -0,0 +1,192 @@ +package modelaccess + +import ( + "context" + "errors" + "net" + "net/http" + "net/http/httptest" + "os/exec" + "path/filepath" + "testing" + "time" + + "github.com/agentserver/agentserver-pkg/internal/modelproxy" + "github.com/agentserver/agentserver-pkg/internal/oauth" + "github.com/agentserver/agentserver-pkg/internal/secrets" + "github.com/agentserver/agentserver-pkg/internal/tokenrefresh" +) + +func TestRunDaemonServesLocalModelProxyAndKeepsServingWithoutRefreshToken(t *testing.T) { + dir := t.TempDir() + sec := secrets.New(filepath.Join(dir, "secrets.json")) + if err := sec.Set(tokenrefresh.AccessTokenKey, "access-from-secret"); err != nil { + t.Fatal(err) + } + + gotAuth := make(chan string, 1) + upstream := httptest.NewServer(http.HandlerFunc(func(w http.ResponseWriter, r *http.Request) { + gotAuth <- r.Header.Get("Authorization") + w.WriteHeader(http.StatusOK) + })) + defer upstream.Close() + + addr := freeTCPAddr(t) + lockPath := filepath.Join(dir, "model-proxy.lock") + ctx, cancel := context.WithCancel(context.Background()) + defer cancel() + errCh := make(chan error, 1) + go func() { + errCh <- RunDaemon(ctx, DaemonOptions{ + Secrets: sec, + OAuth: oauth.AuthCodeConfig{ClientID: "client-x"}, + ProxyAddr: addr, + ProxyUpstreamBaseURL: upstream.URL + "/v1", + LockPath: lockPath, + }) + }() + + waitForProxyHealth(t, "http://"+addr) + lock, err := tokenrefresh.AcquireDaemonLock(lockPath) + if !errors.Is(err, tokenrefresh.ErrDaemonAlreadyRunning) { + if lock != nil { + _ = lock.Close() + } + t.Fatalf("AcquireDaemonLock err=%v, want %v", err, tokenrefresh.ErrDaemonAlreadyRunning) + } + + resp, err := http.Get("http://" + addr + "/v1/models") + if err != nil { + t.Fatalf("proxy request: %v", err) + } + _ = resp.Body.Close() + if resp.StatusCode != http.StatusOK { + t.Fatalf("status = %d, want 200", resp.StatusCode) + } + select { + case got := <-gotAuth: + if got != "Bearer access-from-secret" { + t.Fatalf("Authorization = %q, want bearer access token", got) + } + case <-time.After(time.Second): + t.Fatal("upstream did not receive proxied request") + } + select { + case err := <-errCh: + t.Fatalf("RunDaemon returned before cancellation: %v", err) + default: + } + + cancel() + select { + case err := <-errCh: + if err != nil { + t.Fatalf("RunDaemon returned %v", err) + } + case <-time.After(time.Second): + t.Fatal("RunDaemon did not stop after context cancellation") + } +} + +func TestEnsureDaemonStartsAgentserverModelProxyDaemonWhenHealthMissing(t *testing.T) { + var started bool + var gotArgs []string + healthy := false + err := EnsureDaemon(context.Background(), EnsureDaemonOptions{ + ExePath: "/opt/agentserver/agentserver", + ProxyBaseURL: "http://127.0.0.1:1", + HealthCheck: func(context.Context, string) bool { + return healthy + }, + StartProcess: func(cmd *exec.Cmd) error { + started = true + gotArgs = append([]string(nil), cmd.Args...) + healthy = true + return nil + }, + }) + if err != nil { + t.Fatalf("EnsureDaemon returned %v", err) + } + if !started { + t.Fatal("StartProcess was not called") + } + wantArgs := []string{"/opt/agentserver/agentserver", "model-proxy-daemon"} + if len(gotArgs) != len(wantArgs) { + t.Fatalf("cmd args = %#v, want %#v", gotArgs, wantArgs) + } + for i := range wantArgs { + if gotArgs[i] != wantArgs[i] { + t.Fatalf("cmd args = %#v, want %#v", gotArgs, wantArgs) + } + } +} + +func TestEnsureDaemonReusesHealthyExistingProxy(t *testing.T) { + started := false + err := EnsureDaemon(context.Background(), EnsureDaemonOptions{ + ExePath: "/opt/agentserver/agentserver", + ProxyBaseURL: "http://127.0.0.1:1", + HealthCheck: func(context.Context, string) bool { + return true + }, + StartProcess: func(*exec.Cmd) error { + started = true + return nil + }, + }) + if err != nil { + t.Fatalf("EnsureDaemon returned %v", err) + } + if started { + t.Fatal("StartProcess was called for healthy proxy") + } +} + +func TestEnsureDaemonReportsProxyUnavailableWhenPortOccupied(t *testing.T) { + err := EnsureDaemon(context.Background(), EnsureDaemonOptions{ + ExePath: "/opt/agentserver/agentserver", + ProxyBaseURL: "http://127.0.0.1:1", + HealthCheck: func(context.Context, string) bool { + return false + }, + StartProcess: func(*exec.Cmd) error { + return errors.New("listen tcp 127.0.0.1:53452: bind: address already in use") + }, + }) + if !errors.Is(err, ErrProxyUnavailable) { + t.Fatalf("EnsureDaemon err=%v, want %v", err, ErrProxyUnavailable) + } +} + +func freeTCPAddr(t *testing.T) string { + t.Helper() + ln, err := net.Listen("tcp", "127.0.0.1:0") + if err != nil { + t.Fatal(err) + } + addr := ln.Addr().String() + if err := ln.Close(); err != nil { + t.Fatal(err) + } + return addr +} + +func waitForProxyHealth(t *testing.T, baseURL string) { + t.Helper() + deadline := time.Now().Add(2 * time.Second) + for time.Now().Before(deadline) { + if ProxyHealthy(context.Background(), baseURL) { + return + } + resp, err := http.Get(baseURL + modelproxy.HealthPath) + if err == nil { + _ = resp.Body.Close() + if resp.StatusCode == http.StatusNoContent { + return + } + } + time.Sleep(10 * time.Millisecond) + } + t.Fatalf("model proxy did not become healthy at %s", baseURL) +} From 0d5ca5ce8a4d3d2a0bef508c30cb5eb288ab1d2e Mon Sep 17 00:00:00 2001 From: Zishu Yu Date: Sat, 13 Jun 2026 23:18:01 +0800 Subject: [PATCH 06/40] fix(headless): harden model proxy daemon lifecycle --- internal/modelaccess/daemon.go | 84 +++++++++++---- internal/modelaccess/daemon_test.go | 162 ++++++++++++++++++++++++++++ 2 files changed, 228 insertions(+), 18 deletions(-) diff --git a/internal/modelaccess/daemon.go b/internal/modelaccess/daemon.go index e5ef93a..03b806d 100644 --- a/internal/modelaccess/daemon.go +++ b/internal/modelaccess/daemon.go @@ -4,9 +4,11 @@ import ( "context" "errors" "net/http" + "net/url" "os/exec" "path/filepath" "strings" + "sync" "time" "github.com/agentserver/agentserver-pkg/internal/modelproxy" @@ -19,6 +21,15 @@ import ( var ErrProxyUnavailable = errors.New("modelaccess: local model proxy unavailable") +const healthCheckTimeout = 500 * time.Millisecond + +var ( + runTokenRefresh = tokenrefresh.Run + healthHTTPClient = &http.Client{ + Timeout: healthCheckTimeout, + } +) + type DaemonOptions struct { Secrets secrets.Store OAuth oauth.AuthCodeConfig @@ -33,6 +44,8 @@ func RunDaemon(ctx context.Context, opts DaemonOptions) error { if ctx == nil { ctx = context.Background() } + ctx, cancel := context.WithCancel(ctx) + defer cancel() if opts.Secrets == nil { return tokenrefresh.ErrNoSecrets } @@ -47,46 +60,61 @@ func RunDaemon(ctx context.Context, opts DaemonOptions) error { defer lock.Close() } - proxyErr := make(chan error, 1) + type daemonResult struct { + name string + err error + } + resultCh := make(chan daemonResult, 2) + var wg sync.WaitGroup + finish := func(err error) error { + cancel() + wg.Wait() + return err + } + + wg.Add(1) go func() { - proxyErr <- modelproxy.ListenAndServe(ctx, modelproxy.ServerOptions{ + defer wg.Done() + resultCh <- daemonResult{name: "proxy", err: modelproxy.ListenAndServe(ctx, modelproxy.ServerOptions{ Addr: opts.ProxyAddr, Secrets: opts.Secrets, UpstreamBaseURL: opts.ProxyUpstreamBaseURL, - }) + })} }() - refreshErr := make(chan error, 1) + wg.Add(1) go func() { - refreshErr <- tokenrefresh.Run(ctx, tokenrefresh.Options{ + defer wg.Done() + resultCh <- daemonResult{name: "refresh", err: runTokenRefresh(ctx, tokenrefresh.Options{ Secrets: opts.Secrets, OAuth: opts.OAuth, Refresh: opts.Refresh, Logf: opts.Logf, - }) + })} }() for { select { - case err := <-proxyErr: - if ctx.Err() != nil { - return nil + case result := <-resultCh: + if result.name == "proxy" { + if ctx.Err() != nil { + return finish(nil) + } + return finish(result.err) } - return err - case err := <-refreshErr: + err := result.err if err == nil || errors.Is(err, context.Canceled) { - return nil + return finish(nil) } if errors.Is(err, tokenrefresh.ErrNoRefreshToken) { if opts.Logf != nil { opts.Logf("token refresh disabled: %v", err) } - refreshErr = nil continue } - return err + return finish(err) case <-ctx.Done(): - return nil + return finish(nil) } } } @@ -147,12 +175,17 @@ func ProxyHealthy(ctx context.Context, baseURL string) bool { if ctx == nil { ctx = context.Background() } - baseURL = strings.TrimRight(baseURL, "/") - req, err := http.NewRequestWithContext(ctx, http.MethodGet, baseURL+modelproxy.HealthPath, nil) + ctx, cancel := context.WithTimeout(ctx, healthCheckTimeout) + defer cancel() + healthURL, err := healthURL(baseURL) + if err != nil { + return false + } + req, err := http.NewRequestWithContext(ctx, http.MethodGet, healthURL, nil) if err != nil { return false } - resp, err := http.DefaultClient.Do(req) + resp, err := healthHTTPClient.Do(req) if err != nil { return false } @@ -163,3 +196,18 @@ func ProxyHealthy(ctx context.Context, baseURL string) bool { func DefaultLockPath(installRoot string) string { return filepath.Join(installRoot, "token-refresher.lock") } + +func healthURL(baseURL string) (string, error) { + u, err := url.Parse(strings.TrimRight(baseURL, "/")) + if err != nil { + return "", err + } + if u.Scheme == "" || u.Host == "" { + return "", errors.New("modelaccess: proxy base URL must include scheme and host") + } + u.Path = modelproxy.HealthPath + u.RawPath = "" + u.RawQuery = "" + u.Fragment = "" + return u.String(), nil +} diff --git a/internal/modelaccess/daemon_test.go b/internal/modelaccess/daemon_test.go index 0e85f8e..af2686f 100644 --- a/internal/modelaccess/daemon_test.go +++ b/internal/modelaccess/daemon_test.go @@ -88,6 +88,91 @@ func TestRunDaemonServesLocalModelProxyAndKeepsServingWithoutRefreshToken(t *tes } } +func TestRunDaemonStopsRefreshWhenProxyStartupFails(t *testing.T) { + dir := t.TempDir() + sec := secrets.New(filepath.Join(dir, "secrets.json")) + if err := sec.Set(tokenrefresh.AccessTokenKey, "access-from-secret"); err != nil { + t.Fatal(err) + } + + ln, err := net.Listen("tcp", "127.0.0.1:0") + if err != nil { + t.Fatal(err) + } + defer ln.Close() + + refreshStarted := make(chan struct{}) + refreshStopped := make(chan struct{}) + overrideRunTokenRefresh(t, func(ctx context.Context, _ tokenrefresh.Options) error { + close(refreshStarted) + <-ctx.Done() + close(refreshStopped) + return ctx.Err() + }) + ctx, cancel := context.WithCancel(context.Background()) + defer cancel() + err = RunDaemon(ctx, DaemonOptions{ + Secrets: sec, + OAuth: oauth.AuthCodeConfig{ClientID: "client-x"}, + ProxyAddr: ln.Addr().String(), + }) + if err == nil { + t.Fatal("RunDaemon returned nil, want proxy startup error") + } + select { + case <-refreshStarted: + default: + cancel() + t.Fatal("refresh did not start") + } + select { + case <-refreshStopped: + case <-time.After(time.Second): + cancel() + t.Fatal("refresh was still running after proxy startup failure") + } +} + +func TestRunDaemonStopsProxyWhenRefreshReturnsHardError(t *testing.T) { + dir := t.TempDir() + sec := secrets.New(filepath.Join(dir, "secrets.json")) + if err := sec.Set(tokenrefresh.AccessTokenKey, "access-from-secret"); err != nil { + t.Fatal(err) + } + + addr := freeTCPAddr(t) + refreshErr := errors.New("refresh failed hard") + overrideRunTokenRefresh(t, func(ctx context.Context, _ tokenrefresh.Options) error { + deadline := time.Now().Add(2 * time.Second) + for time.Now().Before(deadline) { + if ProxyHealthy(ctx, "http://"+addr) { + return refreshErr + } + time.Sleep(10 * time.Millisecond) + } + return errors.New("proxy never became healthy") + }) + ctx, cancel := context.WithCancel(context.Background()) + defer cancel() + err := RunDaemon(ctx, DaemonOptions{ + Secrets: sec, + OAuth: oauth.AuthCodeConfig{ClientID: "client-x"}, + ProxyAddr: addr, + }) + if !errors.Is(err, refreshErr) { + t.Fatalf("RunDaemon err=%v, want %v", err, refreshErr) + } + deadline := time.Now().Add(time.Second) + for time.Now().Before(deadline) { + if !ProxyHealthy(context.Background(), "http://"+addr) { + return + } + time.Sleep(10 * time.Millisecond) + } + cancel() + t.Fatal("proxy was still healthy after refresh hard error") +} + func TestEnsureDaemonStartsAgentserverModelProxyDaemonWhenHealthMissing(t *testing.T) { var started bool var gotArgs []string @@ -159,6 +244,74 @@ func TestEnsureDaemonReportsProxyUnavailableWhenPortOccupied(t *testing.T) { } } +func TestProxyHealthyReturnsFalseWhenServerDoesNotRespond(t *testing.T) { + ln, err := net.Listen("tcp", "127.0.0.1:0") + if err != nil { + t.Fatal(err) + } + defer ln.Close() + + connCh := make(chan net.Conn, 1) + acceptErr := make(chan error, 1) + go func() { + conn, err := ln.Accept() + if err != nil { + acceptErr <- err + return + } + connCh <- conn + }() + + done := make(chan bool, 1) + go func() { + done <- ProxyHealthy(context.Background(), "http://"+ln.Addr().String()) + }() + + var conn net.Conn + select { + case conn = <-connCh: + defer conn.Close() + case err := <-acceptErr: + t.Fatalf("Accept: %v", err) + case <-time.After(time.Second): + t.Fatal("ProxyHealthy did not connect to listener") + } + + select { + case healthy := <-done: + if healthy { + t.Fatal("ProxyHealthy returned true for non-responding server") + } + case <-time.After(time.Second): + t.Fatal("ProxyHealthy did not time out") + } +} + +func TestProxyHealthyChecksRootHealthPathWhenBaseURLHasPath(t *testing.T) { + gotPath := make(chan string, 1) + server := httptest.NewServer(http.HandlerFunc(func(w http.ResponseWriter, r *http.Request) { + gotPath <- r.URL.Path + if r.URL.Path != modelproxy.HealthPath { + http.NotFound(w, r) + return + } + w.WriteHeader(http.StatusNoContent) + })) + defer server.Close() + + if !ProxyHealthy(context.Background(), server.URL+"/v1") { + t.Fatal("ProxyHealthy returned false") + } + select { + case got := <-gotPath: + if got != modelproxy.HealthPath { + t.Fatalf("health path = %q, want %q", got, modelproxy.HealthPath) + } + case <-time.After(time.Second): + t.Fatal("server did not receive health check") + } +} + func freeTCPAddr(t *testing.T) string { t.Helper() ln, err := net.Listen("tcp", "127.0.0.1:0") @@ -190,3 +343,12 @@ func waitForProxyHealth(t *testing.T, baseURL string) { } t.Fatalf("model proxy did not become healthy at %s", baseURL) } + +func overrideRunTokenRefresh(t *testing.T, fn func(context.Context, tokenrefresh.Options) error) { + t.Helper() + original := runTokenRefresh + runTokenRefresh = fn + t.Cleanup(func() { + runTokenRefresh = original + }) +} From 0208328df71962ea9da6eeaccf5b285ddc0981fb Mon Sep 17 00:00:00 2001 From: Zishu Yu Date: Sat, 13 Jun 2026 23:24:17 +0800 Subject: [PATCH 07/40] fix(headless): treat daemon context deadlines as shutdown --- internal/modelaccess/daemon.go | 12 +++++++++- internal/modelaccess/daemon_test.go | 35 +++++++++++++++++++++++++++++ 2 files changed, 46 insertions(+), 1 deletion(-) diff --git a/internal/modelaccess/daemon.go b/internal/modelaccess/daemon.go index 03b806d..1c0959b 100644 --- a/internal/modelaccess/daemon.go +++ b/internal/modelaccess/daemon.go @@ -103,7 +103,7 @@ func RunDaemon(ctx context.Context, opts DaemonOptions) error { return finish(result.err) } err := result.err - if err == nil || errors.Is(err, context.Canceled) { + if cleanDaemonShutdown(ctx, err) { return finish(nil) } if errors.Is(err, tokenrefresh.ErrNoRefreshToken) { @@ -197,6 +197,16 @@ func DefaultLockPath(installRoot string) string { return filepath.Join(installRoot, "token-refresher.lock") } +func cleanDaemonShutdown(ctx context.Context, err error) bool { + if err == nil || errors.Is(err, context.Canceled) { + return true + } + if ctxErr := ctx.Err(); ctxErr != nil && errors.Is(err, ctxErr) { + return true + } + return false +} + func healthURL(baseURL string) (string, error) { u, err := url.Parse(strings.TrimRight(baseURL, "/")) if err != nil { diff --git a/internal/modelaccess/daemon_test.go b/internal/modelaccess/daemon_test.go index af2686f..3e35858 100644 --- a/internal/modelaccess/daemon_test.go +++ b/internal/modelaccess/daemon_test.go @@ -173,6 +173,41 @@ func TestRunDaemonStopsProxyWhenRefreshReturnsHardError(t *testing.T) { t.Fatal("proxy was still healthy after refresh hard error") } +func TestRunDaemonReturnsNilWhenDeadlineExpiresFromRefreshBranch(t *testing.T) { + dir := t.TempDir() + sec := secrets.New(filepath.Join(dir, "secrets.json")) + if err := sec.Set(tokenrefresh.AccessTokenKey, "access-from-secret"); err != nil { + t.Fatal(err) + } + + overrideRunTokenRefresh(t, func(ctx context.Context, _ tokenrefresh.Options) error { + <-ctx.Done() + return ctx.Err() + }) + for i := 0; i < 100; i++ { + ctx, cancel := context.WithTimeout(context.Background(), time.Nanosecond) + time.Sleep(time.Millisecond) + err := RunDaemon(ctx, DaemonOptions{ + Secrets: sec, + OAuth: oauth.AuthCodeConfig{ClientID: "client-x"}, + ProxyAddr: freeTCPAddr(t), + }) + cancel() + if err != nil { + t.Fatalf("RunDaemon err=%v on iteration %d, want nil", err, i) + } + } +} + +func TestCleanDaemonShutdownRecognizesActiveContextDeadline(t *testing.T) { + ctx, cancel := context.WithTimeout(context.Background(), time.Nanosecond) + defer cancel() + <-ctx.Done() + if !cleanDaemonShutdown(ctx, context.DeadlineExceeded) { + t.Fatal("context deadline exceeded was not treated as clean shutdown") + } +} + func TestEnsureDaemonStartsAgentserverModelProxyDaemonWhenHealthMissing(t *testing.T) { var started bool var gotArgs []string From e3fcbd8775a7e2f5bf658fd26ebf6efe75de992c Mon Sep 17 00:00:00 2001 From: Zishu Yu Date: Sat, 13 Jun 2026 23:33:56 +0800 Subject: [PATCH 08/40] fix(headless): detach proxy daemon from caller context --- internal/modelaccess/daemon.go | 21 +++++++++-- internal/modelaccess/daemon_test.go | 54 +++++++++++++++++++++++++++++ 2 files changed, 73 insertions(+), 2 deletions(-) diff --git a/internal/modelaccess/daemon.go b/internal/modelaccess/daemon.go index 1c0959b..42ace94 100644 --- a/internal/modelaccess/daemon.go +++ b/internal/modelaccess/daemon.go @@ -25,6 +25,7 @@ const healthCheckTimeout = 500 * time.Millisecond var ( runTokenRefresh = tokenrefresh.Run + newDaemonCommand = exec.Command healthHTTPClient = &http.Client{ Timeout: healthCheckTimeout, } @@ -141,14 +142,20 @@ func EnsureDaemon(ctx context.Context, opts EnsureDaemonOptions) error { if healthCheck(ctx, baseURL) { return nil } - cmd := exec.CommandContext(ctx, opts.ExePath, "model-proxy-daemon") + if opts.ExePath == "" { + return errors.New("agentserver executable path required") + } + if err := ctx.Err(); err != nil { + return err + } + cmd := newDaemonCommand(opts.ExePath, "model-proxy-daemon") cmd.Stdin = nil cmd.Stdout = nil cmd.Stderr = nil process.HideWindow(cmd) startProcess := opts.StartProcess if startProcess == nil { - startProcess = func(cmd *exec.Cmd) error { return cmd.Start() } + startProcess = startDetachedProcess } if err := startProcess(cmd); err != nil { if strings.Contains(strings.ToLower(err.Error()), "address already in use") { @@ -171,6 +178,16 @@ func EnsureDaemon(ctx context.Context, opts EnsureDaemonOptions) error { return ErrProxyUnavailable } +func startDetachedProcess(cmd *exec.Cmd) error { + if err := cmd.Start(); err != nil { + return err + } + if cmd.Process != nil { + return cmd.Process.Release() + } + return nil +} + func ProxyHealthy(ctx context.Context, baseURL string) bool { if ctx == nil { ctx = context.Background() diff --git a/internal/modelaccess/daemon_test.go b/internal/modelaccess/daemon_test.go index 3e35858..0b46945 100644 --- a/internal/modelaccess/daemon_test.go +++ b/internal/modelaccess/daemon_test.go @@ -242,6 +242,51 @@ func TestEnsureDaemonStartsAgentserverModelProxyDaemonWhenHealthMissing(t *testi } } +func TestEnsureDaemonRequiresExePath(t *testing.T) { + err := EnsureDaemon(context.Background(), EnsureDaemonOptions{ + ProxyBaseURL: "http://127.0.0.1:1", + HealthCheck: func(context.Context, string) bool { + return false + }, + StartProcess: func(*exec.Cmd) error { + t.Fatal("StartProcess should not be called without ExePath") + return nil + }, + }) + if err == nil || err.Error() != "agentserver executable path required" { + t.Fatalf("EnsureDaemon err=%v, want agentserver executable path required", err) + } +} + +func TestEnsureDaemonBuildsDetachedCommandNotBoundToCallerContext(t *testing.T) { + parentCtx, cancel := context.WithCancel(context.Background()) + defer cancel() + calledCommandFactory := false + overrideDaemonCommand(t, func(exePath string, args ...string) *exec.Cmd { + calledCommandFactory = true + return exec.Command(exePath, args...) + }) + + healthy := false + err := EnsureDaemon(parentCtx, EnsureDaemonOptions{ + ExePath: "/opt/agentserver/agentserver", + ProxyBaseURL: "http://127.0.0.1:1", + HealthCheck: func(context.Context, string) bool { + return healthy + }, + StartProcess: func(*exec.Cmd) error { + healthy = true + return nil + }, + }) + if err != nil { + t.Fatalf("EnsureDaemon returned %v", err) + } + if !calledCommandFactory { + t.Fatal("daemon command factory was not called") + } +} + func TestEnsureDaemonReusesHealthyExistingProxy(t *testing.T) { started := false err := EnsureDaemon(context.Background(), EnsureDaemonOptions{ @@ -387,3 +432,12 @@ func overrideRunTokenRefresh(t *testing.T, fn func(context.Context, tokenrefresh runTokenRefresh = original }) } + +func overrideDaemonCommand(t *testing.T, fn func(string, ...string) *exec.Cmd) { + t.Helper() + original := newDaemonCommand + newDaemonCommand = fn + t.Cleanup(func() { + newDaemonCommand = original + }) +} From 7f9ee1b08990411b7ce7406526d8cbead135a485 Mon Sep 17 00:00:00 2001 From: Zishu Yu Date: Sat, 13 Jun 2026 23:44:36 +0800 Subject: [PATCH 09/40] feat(headless): enforce modelserver access rule --- internal/modelaccess/access.go | 180 +++++++++++++++++ internal/modelaccess/access_test.go | 297 ++++++++++++++++++++++++++++ 2 files changed, 477 insertions(+) create mode 100644 internal/modelaccess/access.go create mode 100644 internal/modelaccess/access_test.go diff --git a/internal/modelaccess/access.go b/internal/modelaccess/access.go new file mode 100644 index 0000000..5278b1f --- /dev/null +++ b/internal/modelaccess/access.go @@ -0,0 +1,180 @@ +package modelaccess + +import ( + "context" + "errors" + "os" + "time" + + "github.com/agentserver/agentserver-pkg/internal/codex" + "github.com/agentserver/agentserver-pkg/internal/env" + "github.com/agentserver/agentserver-pkg/internal/modelproxy" + "github.com/agentserver/agentserver-pkg/internal/modelserver" + "github.com/agentserver/agentserver-pkg/internal/oauth" + "github.com/agentserver/agentserver-pkg/internal/secrets" + "github.com/agentserver/agentserver-pkg/internal/tokenrefresh" +) + +type Mode string + +const ( + ModeDirectAPIKey Mode = "direct_api_key" + ModeLocalProxy Mode = "local_proxy" +) + +type Result struct { + Mode Mode +} + +type EnsureOptions struct { + CodexConfigPath string + Secrets secrets.Store + + DeviceConfig oauth.Config + AuthCodeConfig oauth.AuthCodeConfig + + Env func(string) string + SetEnv func(string, string) error + PersistEnv func(string, string) error + Now func() time.Time + + RequestDeviceCode func(context.Context, oauth.Config) (oauth.DeviceCodeChallenge, error) + PrintChallenge func(string, oauth.DeviceCodeChallenge) + PollToken func(context.Context, oauth.Config, oauth.DeviceCodeChallenge) (oauth.Token, error) + Refresh func(context.Context, oauth.AuthCodeConfig, string) (oauth.Token, error) + StartDaemon func(context.Context) error +} + +func Ensure(ctx context.Context, opts EnsureOptions) (Result, error) { + if ctx == nil { + ctx = context.Background() + } + opts = defaultEnsureOptions(opts) + if opts.Secrets == nil { + return Result{}, tokenrefresh.ErrNoSecrets + } + + if opts.Env(tokenrefresh.OpenAIAPIKeyEnv) != "" { + if err := codex.UpdateConfig(opts.CodexConfigPath, codex.ModelserverSettings()); err != nil { + return Result{}, err + } + return Result{Mode: ModeDirectAPIKey}, nil + } + + if err := ensureProxyCredentials(ctx, opts); err != nil { + return Result{}, err + } + if err := codex.UpdateConfig(opts.CodexConfigPath, codex.ModelserverProxySettings(modelproxy.DefaultBaseURL)); err != nil { + return Result{}, err + } + if err := opts.SetEnv(codex.LocalProxyAPIKeyEnv, codex.LocalProxyAPIKeyValue); err != nil { + return Result{}, err + } + if err := opts.PersistEnv(codex.LocalProxyAPIKeyEnv, codex.LocalProxyAPIKeyValue); err != nil { + return Result{}, err + } + if opts.StartDaemon != nil { + if err := opts.StartDaemon(ctx); err != nil { + return Result{}, err + } + } + return Result{Mode: ModeLocalProxy}, nil +} + +func ensureProxyCredentials(ctx context.Context, opts EnsureOptions) error { + reauth, err := opts.Secrets.Get(tokenrefresh.ReauthRequiredKey) + if err != nil && !errors.Is(err, secrets.ErrNotFound) { + return err + } + if reauth == "true" { + return runDeviceLogin(ctx, opts) + } + + refreshToken, err := opts.Secrets.Get(tokenrefresh.RefreshTokenKey) + if err != nil { + if errors.Is(err, secrets.ErrNotFound) { + return runDeviceLogin(ctx, opts) + } + return err + } + if refreshToken == "" { + return runDeviceLogin(ctx, opts) + } + + accessToken, err := opts.Secrets.Get(tokenrefresh.AccessTokenKey) + if err != nil { + if errors.Is(err, secrets.ErrNotFound) { + err := refreshOnce(ctx, opts) + if tokenrefresh.ReauthRequired(err) { + return runDeviceLogin(ctx, opts) + } + return err + } + return err + } + if accessToken == "" { + err := refreshOnce(ctx, opts) + if tokenrefresh.ReauthRequired(err) { + return runDeviceLogin(ctx, opts) + } + return err + } + return nil +} + +func refreshOnce(ctx context.Context, opts EnsureOptions) error { + _, err := tokenrefresh.RefreshOnce(ctx, tokenrefresh.Options{ + Secrets: opts.Secrets, + OAuth: opts.AuthCodeConfig, + Refresh: opts.Refresh, + Now: opts.Now, + }) + return err +} + +func runDeviceLogin(ctx context.Context, opts EnsureOptions) error { + ch, err := opts.RequestDeviceCode(ctx, opts.DeviceConfig) + if err != nil { + return err + } + if opts.PrintChallenge != nil { + opts.PrintChallenge("Code 登录", ch) + } + tok, err := opts.PollToken(ctx, opts.DeviceConfig, ch) + if err != nil { + return err + } + _, err = tokenrefresh.StoreToken(opts.Secrets, tok, opts.Now(), "") + return err +} + +func defaultEnsureOptions(opts EnsureOptions) EnsureOptions { + if opts.DeviceConfig.ClientID == "" { + opts.DeviceConfig = modelserver.DeviceConfig("") + } + if opts.AuthCodeConfig.ClientID == "" { + opts.AuthCodeConfig = modelserver.OAuthConfig() + } + if opts.Env == nil { + opts.Env = os.Getenv + } + if opts.SetEnv == nil { + opts.SetEnv = os.Setenv + } + if opts.PersistEnv == nil { + opts.PersistEnv = env.PersistUserEnv + } + if opts.Now == nil { + opts.Now = func() time.Time { return time.Now().UTC() } + } + if opts.RequestDeviceCode == nil { + opts.RequestDeviceCode = oauth.RequestDeviceCode + } + if opts.PollToken == nil { + opts.PollToken = oauth.PollToken + } + if opts.Refresh == nil { + opts.Refresh = oauth.RefreshToken + } + return opts +} diff --git a/internal/modelaccess/access_test.go b/internal/modelaccess/access_test.go new file mode 100644 index 0000000..571cb03 --- /dev/null +++ b/internal/modelaccess/access_test.go @@ -0,0 +1,297 @@ +package modelaccess + +import ( + "context" + "errors" + "os" + "path/filepath" + "strings" + "testing" + "time" + + "github.com/agentserver/agentserver-pkg/internal/codex" + "github.com/agentserver/agentserver-pkg/internal/modelproxy" + "github.com/agentserver/agentserver-pkg/internal/oauth" + "github.com/agentserver/agentserver-pkg/internal/secrets" + "github.com/agentserver/agentserver-pkg/internal/tokenrefresh" +) + +func TestEnsurePrefersLongLivedAPIKey(t *testing.T) { + tmp := t.TempDir() + sec := secrets.New(filepath.Join(tmp, "secrets.json")) + var daemonStarted bool + + result, err := Ensure(context.Background(), EnsureOptions{ + CodexConfigPath: filepath.Join(tmp, "codex", "config.toml"), + Secrets: sec, + Env: func(key string) string { + if key == tokenrefresh.OpenAIAPIKeyEnv { + return "sk-test" + } + return "" + }, + StartDaemon: func(context.Context) error { + daemonStarted = true + return nil + }, + }) + if err != nil { + t.Fatalf("Ensure() error = %v", err) + } + if result.Mode != ModeDirectAPIKey { + t.Fatalf("Mode = %q, want %q", result.Mode, ModeDirectAPIKey) + } + if daemonStarted { + t.Fatal("StartDaemon called in direct API key mode") + } + assertConfigContains(t, filepath.Join(tmp, "codex", "config.toml"), + `model_provider = "modelserver"`, + `base_url = "https://code.ai.cs.ac.cn/v1"`, + `env_key = "OPENAI_API_KEY"`, + ) +} + +func TestEnsureProxyModeRunsDeviceLoginWhenRefreshMissing(t *testing.T) { + tmp := t.TempDir() + sec := secrets.New(filepath.Join(tmp, "secrets.json")) + challenge := oauth.DeviceCodeChallenge{ + DeviceCode: "device-code", + UserCode: "ABCD-EFGH", + VerificationURI: "https://example.test/device", + ExpiresIn: 600, + Interval: 1, + } + var requestedDeviceCode, polledToken, printedChallenge, daemonStarted bool + var printedTitle string + + result, err := Ensure(context.Background(), EnsureOptions{ + CodexConfigPath: filepath.Join(tmp, "codex", "config.toml"), + Secrets: sec, + Env: func(string) string { return "" }, + RequestDeviceCode: func(context.Context, oauth.Config) (oauth.DeviceCodeChallenge, error) { + requestedDeviceCode = true + return challenge, nil + }, + PrintChallenge: func(title string, ch oauth.DeviceCodeChallenge) { + printedChallenge = true + printedTitle = title + if ch.DeviceCode != challenge.DeviceCode { + t.Fatalf("printed challenge DeviceCode = %q, want %q", ch.DeviceCode, challenge.DeviceCode) + } + }, + PollToken: func(context.Context, oauth.Config, oauth.DeviceCodeChallenge) (oauth.Token, error) { + polledToken = true + return oauth.Token{AccessToken: "access", RefreshToken: "refresh", ExpiresIn: 3600}, nil + }, + StartDaemon: func(context.Context) error { + daemonStarted = true + return nil + }, + Now: fixedNow, + }) + if err != nil { + t.Fatalf("Ensure() error = %v", err) + } + if result.Mode != ModeLocalProxy { + t.Fatalf("Mode = %q, want %q", result.Mode, ModeLocalProxy) + } + if !requestedDeviceCode || !printedChallenge || !polledToken || !daemonStarted { + t.Fatalf("calls: request=%v print=%v poll=%v daemon=%v", requestedDeviceCode, printedChallenge, polledToken, daemonStarted) + } + if printedTitle != "Code 登录" { + t.Fatalf("PrintChallenge title = %q, want %q", printedTitle, "Code 登录") + } + assertSecret(t, sec, tokenrefresh.AccessTokenKey, "access") + assertSecret(t, sec, tokenrefresh.RefreshTokenKey, "refresh") +} + +func TestEnsureProxyModeUsesExistingRefreshTokenWithoutPrompt(t *testing.T) { + tmp := t.TempDir() + sec := secrets.New(filepath.Join(tmp, "secrets.json")) + mustSetSecret(t, sec, tokenrefresh.AccessTokenKey, "existing-access") + mustSetSecret(t, sec, tokenrefresh.RefreshTokenKey, "existing-refresh") + var daemonStarted bool + + result, err := Ensure(context.Background(), EnsureOptions{ + CodexConfigPath: filepath.Join(tmp, "codex", "config.toml"), + Secrets: sec, + Env: func(string) string { return "" }, + RequestDeviceCode: func(context.Context, oauth.Config) (oauth.DeviceCodeChallenge, error) { + t.Fatal("RequestDeviceCode called") + return oauth.DeviceCodeChallenge{}, nil + }, + PollToken: func(context.Context, oauth.Config, oauth.DeviceCodeChallenge) (oauth.Token, error) { + t.Fatal("PollToken called") + return oauth.Token{}, nil + }, + StartDaemon: func(context.Context) error { + daemonStarted = true + return nil + }, + }) + if err != nil { + t.Fatalf("Ensure() error = %v", err) + } + if result.Mode != ModeLocalProxy { + t.Fatalf("Mode = %q, want %q", result.Mode, ModeLocalProxy) + } + if !daemonStarted { + t.Fatal("StartDaemon was not called") + } + assertSecret(t, sec, tokenrefresh.AccessTokenKey, "existing-access") + assertSecret(t, sec, tokenrefresh.RefreshTokenKey, "existing-refresh") +} + +func TestEnsureProxyModeRefreshesWhenAccessTokenMissing(t *testing.T) { + tmp := t.TempDir() + sec := secrets.New(filepath.Join(tmp, "secrets.json")) + mustSetSecret(t, sec, tokenrefresh.RefreshTokenKey, "existing-refresh") + var refreshed, daemonStarted bool + + result, err := Ensure(context.Background(), EnsureOptions{ + CodexConfigPath: filepath.Join(tmp, "codex", "config.toml"), + Secrets: sec, + Env: func(string) string { return "" }, + Refresh: func(_ context.Context, _ oauth.AuthCodeConfig, refreshToken string) (oauth.Token, error) { + refreshed = true + if refreshToken != "existing-refresh" { + t.Fatalf("refresh token = %q, want %q", refreshToken, "existing-refresh") + } + return oauth.Token{ + AccessToken: "refreshed-access", + ExpiresIn: 3600, + }, nil + }, + RequestDeviceCode: func(context.Context, oauth.Config) (oauth.DeviceCodeChallenge, error) { + t.Fatal("RequestDeviceCode called") + return oauth.DeviceCodeChallenge{}, nil + }, + StartDaemon: func(context.Context) error { + daemonStarted = true + return nil + }, + Now: fixedNow, + }) + if err != nil { + t.Fatalf("Ensure() error = %v", err) + } + if result.Mode != ModeLocalProxy { + t.Fatalf("Mode = %q, want %q", result.Mode, ModeLocalProxy) + } + if !refreshed || !daemonStarted { + t.Fatalf("calls: refresh=%v daemon=%v", refreshed, daemonStarted) + } + assertSecret(t, sec, tokenrefresh.AccessTokenKey, "refreshed-access") + assertSecret(t, sec, tokenrefresh.RefreshTokenKey, "existing-refresh") +} + +func TestEnsureRerunsLoginWhenReauthFlagSet(t *testing.T) { + tmp := t.TempDir() + sec := secrets.New(filepath.Join(tmp, "secrets.json")) + mustSetSecret(t, sec, tokenrefresh.AccessTokenKey, "stale-access") + mustSetSecret(t, sec, tokenrefresh.RefreshTokenKey, "stale-refresh") + mustSetSecret(t, sec, tokenrefresh.ReauthRequiredKey, "true") + var requestedDeviceCode, daemonStarted bool + + _, err := Ensure(context.Background(), EnsureOptions{ + CodexConfigPath: filepath.Join(tmp, "codex", "config.toml"), + Secrets: sec, + Env: func(string) string { return "" }, + RequestDeviceCode: func(context.Context, oauth.Config) (oauth.DeviceCodeChallenge, error) { + requestedDeviceCode = true + return oauth.DeviceCodeChallenge{DeviceCode: "device", UserCode: "code", ExpiresIn: 600}, nil + }, + PrintChallenge: func(string, oauth.DeviceCodeChallenge) {}, + PollToken: func(context.Context, oauth.Config, oauth.DeviceCodeChallenge) (oauth.Token, error) { + return oauth.Token{AccessToken: "new-access", RefreshToken: "new-refresh", ExpiresIn: 3600}, nil + }, + StartDaemon: func(context.Context) error { + daemonStarted = true + return nil + }, + Now: fixedNow, + }) + if err != nil { + t.Fatalf("Ensure() error = %v", err) + } + if !requestedDeviceCode || !daemonStarted { + t.Fatalf("calls: request=%v daemon=%v", requestedDeviceCode, daemonStarted) + } + assertSecret(t, sec, tokenrefresh.AccessTokenKey, "new-access") + assertSecret(t, sec, tokenrefresh.RefreshTokenKey, "new-refresh") + if _, err := sec.Get(tokenrefresh.ReauthRequiredKey); !errors.Is(err, secrets.ErrNotFound) { + t.Fatalf("ReauthRequiredKey error = %v, want ErrNotFound", err) + } +} + +func TestProxySettingsWrittenInProxyMode(t *testing.T) { + tmp := t.TempDir() + sec := secrets.New(filepath.Join(tmp, "secrets.json")) + mustSetSecret(t, sec, tokenrefresh.AccessTokenKey, "existing-access") + mustSetSecret(t, sec, tokenrefresh.RefreshTokenKey, "existing-refresh") + var setEnvKey, setEnvValue, persistEnvKey, persistEnvValue string + + _, err := Ensure(context.Background(), EnsureOptions{ + CodexConfigPath: filepath.Join(tmp, "codex", "config.toml"), + Secrets: sec, + Env: func(string) string { return "" }, + SetEnv: func(key, value string) error { + setEnvKey, setEnvValue = key, value + return nil + }, + PersistEnv: func(key, value string) error { + persistEnvKey, persistEnvValue = key, value + return nil + }, + StartDaemon: func(context.Context) error { return nil }, + }) + if err != nil { + t.Fatalf("Ensure() error = %v", err) + } + assertConfigContains(t, filepath.Join(tmp, "codex", "config.toml"), + `base_url = "`+modelproxy.DefaultBaseURL+`"`, + `env_key = "`+codex.LocalProxyAPIKeyEnv+`"`, + ) + if setEnvKey != codex.LocalProxyAPIKeyEnv || setEnvValue != codex.LocalProxyAPIKeyValue { + t.Fatalf("SetEnv = (%q, %q), want (%q, %q)", setEnvKey, setEnvValue, codex.LocalProxyAPIKeyEnv, codex.LocalProxyAPIKeyValue) + } + if persistEnvKey != codex.LocalProxyAPIKeyEnv || persistEnvValue != codex.LocalProxyAPIKeyValue { + t.Fatalf("PersistEnv = (%q, %q), want (%q, %q)", persistEnvKey, persistEnvValue, codex.LocalProxyAPIKeyEnv, codex.LocalProxyAPIKeyValue) + } +} + +func fixedNow() time.Time { + return time.Date(2026, 6, 13, 1, 2, 3, 0, time.UTC) +} + +func mustSetSecret(t *testing.T, sec secrets.Store, key, value string) { + t.Helper() + if err := sec.Set(key, value); err != nil { + t.Fatalf("Set(%q) error = %v", key, err) + } +} + +func assertSecret(t *testing.T, sec secrets.Store, key, want string) { + t.Helper() + got, err := sec.Get(key) + if err != nil { + t.Fatalf("Get(%q) error = %v", key, err) + } + if got != want { + t.Fatalf("Get(%q) = %q, want %q", key, got, want) + } +} + +func assertConfigContains(t *testing.T, path string, wants ...string) { + t.Helper() + b, err := os.ReadFile(path) + if err != nil { + t.Fatalf("ReadFile(%q) error = %v", path, err) + } + text := string(b) + for _, want := range wants { + if !strings.Contains(text, want) { + t.Fatalf("config missing %q:\n%s", want, text) + } + } +} From 72cb17abe7f507cbf808ef9d2c79f61f2cd6ebd8 Mon Sep 17 00:00:00 2001 From: Zishu Yu Date: Sat, 13 Jun 2026 23:48:28 +0800 Subject: [PATCH 10/40] test(headless): cover model access edge cases --- internal/modelaccess/access_test.go | 110 ++++++++++++++++++++++++++++ 1 file changed, 110 insertions(+) diff --git a/internal/modelaccess/access_test.go b/internal/modelaccess/access_test.go index 571cb03..efd4462 100644 --- a/internal/modelaccess/access_test.go +++ b/internal/modelaccess/access_test.go @@ -3,6 +3,7 @@ package modelaccess import ( "context" "errors" + "fmt" "os" "path/filepath" "strings" @@ -34,6 +35,14 @@ func TestEnsurePrefersLongLivedAPIKey(t *testing.T) { daemonStarted = true return nil }, + RequestDeviceCode: func(context.Context, oauth.Config) (oauth.DeviceCodeChallenge, error) { + t.Fatal("RequestDeviceCode called in direct API key mode") + return oauth.DeviceCodeChallenge{}, nil + }, + PollToken: func(context.Context, oauth.Config, oauth.DeviceCodeChallenge) (oauth.Token, error) { + t.Fatal("PollToken called in direct API key mode") + return oauth.Token{}, nil + }, }) if err != nil { t.Fatalf("Ensure() error = %v", err) @@ -185,6 +194,97 @@ func TestEnsureProxyModeRefreshesWhenAccessTokenMissing(t *testing.T) { assertSecret(t, sec, tokenrefresh.RefreshTokenKey, "existing-refresh") } +func TestEnsureProxyModeRunsDeviceLoginWhenRefreshNeedsReauth(t *testing.T) { + tmp := t.TempDir() + sec := secrets.New(filepath.Join(tmp, "secrets.json")) + mustSetSecret(t, sec, tokenrefresh.RefreshTokenKey, "existing-refresh") + var refreshed, requestedDeviceCode, polledToken, daemonStarted bool + + result, err := Ensure(context.Background(), EnsureOptions{ + CodexConfigPath: filepath.Join(tmp, "codex", "config.toml"), + Secrets: sec, + Env: func(string) string { return "" }, + Refresh: func(_ context.Context, _ oauth.AuthCodeConfig, refreshToken string) (oauth.Token, error) { + refreshed = true + if refreshToken != "existing-refresh" { + t.Fatalf("refresh token = %q, want %q", refreshToken, "existing-refresh") + } + return oauth.Token{}, fmt.Errorf("refresh failed: %w", oauth.ErrInvalidGrant) + }, + RequestDeviceCode: func(context.Context, oauth.Config) (oauth.DeviceCodeChallenge, error) { + requestedDeviceCode = true + return oauth.DeviceCodeChallenge{DeviceCode: "device", UserCode: "code", ExpiresIn: 600}, nil + }, + PrintChallenge: func(string, oauth.DeviceCodeChallenge) {}, + PollToken: func(context.Context, oauth.Config, oauth.DeviceCodeChallenge) (oauth.Token, error) { + polledToken = true + return oauth.Token{AccessToken: "device-access", RefreshToken: "device-refresh", ExpiresIn: 3600}, nil + }, + StartDaemon: func(context.Context) error { + daemonStarted = true + return nil + }, + Now: fixedNow, + }) + if err != nil { + t.Fatalf("Ensure() error = %v", err) + } + if result.Mode != ModeLocalProxy { + t.Fatalf("Mode = %q, want %q", result.Mode, ModeLocalProxy) + } + if !refreshed || !requestedDeviceCode || !polledToken || !daemonStarted { + t.Fatalf("calls: refresh=%v request=%v poll=%v daemon=%v", refreshed, requestedDeviceCode, polledToken, daemonStarted) + } + assertSecret(t, sec, tokenrefresh.AccessTokenKey, "device-access") + assertSecret(t, sec, tokenrefresh.RefreshTokenKey, "device-refresh") + if _, err := sec.Get(tokenrefresh.ReauthRequiredKey); !errors.Is(err, secrets.ErrNotFound) { + t.Fatalf("ReauthRequiredKey error = %v, want ErrNotFound", err) + } +} + +func TestEnsureProxyModePropagatesRefreshErrorWithoutPrompt(t *testing.T) { + tmp := t.TempDir() + sec := secrets.New(filepath.Join(tmp, "secrets.json")) + mustSetSecret(t, sec, tokenrefresh.RefreshTokenKey, "existing-refresh") + refreshErr := errors.New("refresh unavailable") + var refreshed bool + + _, err := Ensure(context.Background(), EnsureOptions{ + CodexConfigPath: filepath.Join(tmp, "codex", "config.toml"), + Secrets: sec, + Env: func(string) string { return "" }, + Refresh: func(_ context.Context, _ oauth.AuthCodeConfig, refreshToken string) (oauth.Token, error) { + refreshed = true + if refreshToken != "existing-refresh" { + t.Fatalf("refresh token = %q, want %q", refreshToken, "existing-refresh") + } + return oauth.Token{}, refreshErr + }, + RequestDeviceCode: func(context.Context, oauth.Config) (oauth.DeviceCodeChallenge, error) { + t.Fatal("RequestDeviceCode called after non-reauth refresh error") + return oauth.DeviceCodeChallenge{}, nil + }, + PollToken: func(context.Context, oauth.Config, oauth.DeviceCodeChallenge) (oauth.Token, error) { + t.Fatal("PollToken called after non-reauth refresh error") + return oauth.Token{}, nil + }, + StartDaemon: func(context.Context) error { + t.Fatal("StartDaemon called after refresh error") + return nil + }, + Now: fixedNow, + }) + if !errors.Is(err, refreshErr) { + t.Fatalf("Ensure() error = %v, want %v", err, refreshErr) + } + if !refreshed { + t.Fatal("Refresh was not called") + } + if _, err := sec.Get(tokenrefresh.AccessTokenKey); !errors.Is(err, secrets.ErrNotFound) { + t.Fatalf("AccessTokenKey error = %v, want ErrNotFound", err) + } +} + func TestEnsureRerunsLoginWhenReauthFlagSet(t *testing.T) { tmp := t.TempDir() sec := secrets.New(filepath.Join(tmp, "secrets.json")) @@ -224,6 +324,16 @@ func TestEnsureRerunsLoginWhenReauthFlagSet(t *testing.T) { } } +func TestEnsureRequiresSecrets(t *testing.T) { + _, err := Ensure(context.Background(), EnsureOptions{ + CodexConfigPath: filepath.Join(t.TempDir(), "codex", "config.toml"), + Env: func(string) string { return "" }, + }) + if !errors.Is(err, tokenrefresh.ErrNoSecrets) { + t.Fatalf("Ensure() error = %v, want %v", err, tokenrefresh.ErrNoSecrets) + } +} + func TestProxySettingsWrittenInProxyMode(t *testing.T) { tmp := t.TempDir() sec := secrets.New(filepath.Join(tmp, "secrets.json")) From 2fed96b948db0dd0c0283dac994930343d325f45 Mon Sep 17 00:00:00 2001 From: Zishu Yu Date: Sat, 13 Jun 2026 23:53:13 +0800 Subject: [PATCH 11/40] fix(headless): allow direct model access without secrets --- internal/modelaccess/access.go | 6 +++--- internal/modelaccess/access_test.go | 26 ++++++++++++++++++++++++++ 2 files changed, 29 insertions(+), 3 deletions(-) diff --git a/internal/modelaccess/access.go b/internal/modelaccess/access.go index 5278b1f..be7e0e3 100644 --- a/internal/modelaccess/access.go +++ b/internal/modelaccess/access.go @@ -50,9 +50,6 @@ func Ensure(ctx context.Context, opts EnsureOptions) (Result, error) { ctx = context.Background() } opts = defaultEnsureOptions(opts) - if opts.Secrets == nil { - return Result{}, tokenrefresh.ErrNoSecrets - } if opts.Env(tokenrefresh.OpenAIAPIKeyEnv) != "" { if err := codex.UpdateConfig(opts.CodexConfigPath, codex.ModelserverSettings()); err != nil { @@ -61,6 +58,9 @@ func Ensure(ctx context.Context, opts EnsureOptions) (Result, error) { return Result{Mode: ModeDirectAPIKey}, nil } + if opts.Secrets == nil { + return Result{}, tokenrefresh.ErrNoSecrets + } if err := ensureProxyCredentials(ctx, opts); err != nil { return Result{}, err } diff --git a/internal/modelaccess/access_test.go b/internal/modelaccess/access_test.go index efd4462..9176241 100644 --- a/internal/modelaccess/access_test.go +++ b/internal/modelaccess/access_test.go @@ -60,6 +60,32 @@ func TestEnsurePrefersLongLivedAPIKey(t *testing.T) { ) } +func TestEnsureDirectAPIKeyModeDoesNotRequireSecrets(t *testing.T) { + tmp := t.TempDir() + configPath := filepath.Join(tmp, "codex", "config.toml") + + result, err := Ensure(context.Background(), EnsureOptions{ + CodexConfigPath: configPath, + Env: func(key string) string { + if key == tokenrefresh.OpenAIAPIKeyEnv { + return "sk-test" + } + return "" + }, + }) + if err != nil { + t.Fatalf("Ensure() error = %v", err) + } + if result.Mode != ModeDirectAPIKey { + t.Fatalf("Mode = %q, want %q", result.Mode, ModeDirectAPIKey) + } + assertConfigContains(t, configPath, + `model_provider = "modelserver"`, + `base_url = "https://code.ai.cs.ac.cn/v1"`, + `env_key = "OPENAI_API_KEY"`, + ) +} + func TestEnsureProxyModeRunsDeviceLoginWhenRefreshMissing(t *testing.T) { tmp := t.TempDir() sec := secrets.New(filepath.Join(tmp, "secrets.json")) From 3616684a3cdadd9ceded4c05e0574fe23a3fb3cc Mon Sep 17 00:00:00 2001 From: Zishu Yu Date: Sun, 14 Jun 2026 00:00:25 +0800 Subject: [PATCH 12/40] fix(headless): avoid proxy config on access setup failure --- internal/modelaccess/access.go | 6 +- internal/modelaccess/access_test.go | 112 +++++++++++++++++++++++++++- 2 files changed, 114 insertions(+), 4 deletions(-) diff --git a/internal/modelaccess/access.go b/internal/modelaccess/access.go index be7e0e3..a9b500f 100644 --- a/internal/modelaccess/access.go +++ b/internal/modelaccess/access.go @@ -64,9 +64,6 @@ func Ensure(ctx context.Context, opts EnsureOptions) (Result, error) { if err := ensureProxyCredentials(ctx, opts); err != nil { return Result{}, err } - if err := codex.UpdateConfig(opts.CodexConfigPath, codex.ModelserverProxySettings(modelproxy.DefaultBaseURL)); err != nil { - return Result{}, err - } if err := opts.SetEnv(codex.LocalProxyAPIKeyEnv, codex.LocalProxyAPIKeyValue); err != nil { return Result{}, err } @@ -78,6 +75,9 @@ func Ensure(ctx context.Context, opts EnsureOptions) (Result, error) { return Result{}, err } } + if err := codex.UpdateConfig(opts.CodexConfigPath, codex.ModelserverProxySettings(modelproxy.DefaultBaseURL)); err != nil { + return Result{}, err + } return Result{Mode: ModeLocalProxy}, nil } diff --git a/internal/modelaccess/access_test.go b/internal/modelaccess/access_test.go index 9176241..e4685ad 100644 --- a/internal/modelaccess/access_test.go +++ b/internal/modelaccess/access_test.go @@ -360,12 +360,97 @@ func TestEnsureRequiresSecrets(t *testing.T) { } } +func TestEnsureProxyModeDoesNotWriteProxyConfigWhenSetEnvFails(t *testing.T) { + tmp := t.TempDir() + configPath := filepath.Join(tmp, "codex", "config.toml") + writeDirectConfig(t, configPath) + sec := secrets.New(filepath.Join(tmp, "secrets.json")) + mustSetSecret(t, sec, tokenrefresh.AccessTokenKey, "existing-access") + mustSetSecret(t, sec, tokenrefresh.RefreshTokenKey, "existing-refresh") + setEnvErr := errors.New("set env failed") + + _, err := Ensure(context.Background(), EnsureOptions{ + CodexConfigPath: configPath, + Secrets: sec, + Env: func(string) string { return "" }, + SetEnv: func(string, string) error { + return setEnvErr + }, + PersistEnv: func(string, string) error { + t.Fatal("PersistEnv called after SetEnv failure") + return nil + }, + StartDaemon: func(context.Context) error { + t.Fatal("StartDaemon called after SetEnv failure") + return nil + }, + }) + if !errors.Is(err, setEnvErr) { + t.Fatalf("Ensure() error = %v, want %v", err, setEnvErr) + } + assertConfigNotContains(t, configPath, `base_url = "`+modelproxy.DefaultBaseURL+`"`) +} + +func TestEnsureProxyModeDoesNotWriteProxyConfigWhenPersistEnvFails(t *testing.T) { + tmp := t.TempDir() + configPath := filepath.Join(tmp, "codex", "config.toml") + writeDirectConfig(t, configPath) + sec := secrets.New(filepath.Join(tmp, "secrets.json")) + mustSetSecret(t, sec, tokenrefresh.AccessTokenKey, "existing-access") + mustSetSecret(t, sec, tokenrefresh.RefreshTokenKey, "existing-refresh") + persistErr := errors.New("persist env failed") + + _, err := Ensure(context.Background(), EnsureOptions{ + CodexConfigPath: configPath, + Secrets: sec, + Env: func(string) string { return "" }, + SetEnv: func(string, string) error { return nil }, + PersistEnv: func(string, string) error { + return persistErr + }, + StartDaemon: func(context.Context) error { + t.Fatal("StartDaemon called after PersistEnv failure") + return nil + }, + }) + if !errors.Is(err, persistErr) { + t.Fatalf("Ensure() error = %v, want %v", err, persistErr) + } + assertConfigNotContains(t, configPath, `base_url = "`+modelproxy.DefaultBaseURL+`"`) +} + +func TestEnsureProxyModeDoesNotWriteProxyConfigWhenStartDaemonFails(t *testing.T) { + tmp := t.TempDir() + configPath := filepath.Join(tmp, "codex", "config.toml") + writeDirectConfig(t, configPath) + sec := secrets.New(filepath.Join(tmp, "secrets.json")) + mustSetSecret(t, sec, tokenrefresh.AccessTokenKey, "existing-access") + mustSetSecret(t, sec, tokenrefresh.RefreshTokenKey, "existing-refresh") + daemonErr := errors.New("daemon failed") + + _, err := Ensure(context.Background(), EnsureOptions{ + CodexConfigPath: configPath, + Secrets: sec, + Env: func(string) string { return "" }, + SetEnv: func(string, string) error { return nil }, + PersistEnv: func(string, string) error { return nil }, + StartDaemon: func(context.Context) error { + return daemonErr + }, + }) + if !errors.Is(err, daemonErr) { + t.Fatalf("Ensure() error = %v, want %v", err, daemonErr) + } + assertConfigNotContains(t, configPath, `base_url = "`+modelproxy.DefaultBaseURL+`"`) +} + func TestProxySettingsWrittenInProxyMode(t *testing.T) { tmp := t.TempDir() sec := secrets.New(filepath.Join(tmp, "secrets.json")) mustSetSecret(t, sec, tokenrefresh.AccessTokenKey, "existing-access") mustSetSecret(t, sec, tokenrefresh.RefreshTokenKey, "existing-refresh") var setEnvKey, setEnvValue, persistEnvKey, persistEnvValue string + var daemonStarted bool _, err := Ensure(context.Background(), EnsureOptions{ CodexConfigPath: filepath.Join(tmp, "codex", "config.toml"), @@ -379,11 +464,17 @@ func TestProxySettingsWrittenInProxyMode(t *testing.T) { persistEnvKey, persistEnvValue = key, value return nil }, - StartDaemon: func(context.Context) error { return nil }, + StartDaemon: func(context.Context) error { + daemonStarted = true + return nil + }, }) if err != nil { t.Fatalf("Ensure() error = %v", err) } + if !daemonStarted { + t.Fatal("StartDaemon was not called") + } assertConfigContains(t, filepath.Join(tmp, "codex", "config.toml"), `base_url = "`+modelproxy.DefaultBaseURL+`"`, `env_key = "`+codex.LocalProxyAPIKeyEnv+`"`, @@ -396,6 +487,13 @@ func TestProxySettingsWrittenInProxyMode(t *testing.T) { } } +func writeDirectConfig(t *testing.T, path string) { + t.Helper() + if err := codex.UpdateConfig(path, codex.ModelserverSettings()); err != nil { + t.Fatalf("UpdateConfig() error = %v", err) + } +} + func fixedNow() time.Time { return time.Date(2026, 6, 13, 1, 2, 3, 0, time.UTC) } @@ -431,3 +529,15 @@ func assertConfigContains(t *testing.T, path string, wants ...string) { } } } + +func assertConfigNotContains(t *testing.T, path, unwanted string) { + t.Helper() + b, err := os.ReadFile(path) + if err != nil { + t.Fatalf("ReadFile(%q) error = %v", path, err) + } + text := string(b) + if strings.Contains(text, unwanted) { + t.Fatalf("config contains %q:\n%s", unwanted, text) + } +} From 366bd1d465b50eddf63cd69002960beb36e46ebc Mon Sep 17 00:00:00 2001 From: Zishu Yu Date: Sun, 14 Jun 2026 00:11:20 +0800 Subject: [PATCH 13/40] feat(headless): resolve Linux runtime assets --- internal/headless/runtime.go | 86 ++++++++++ internal/headless/runtime_test.go | 161 ++++++++++++++++++ .../linux/codex-manifest-linux-amd64.json | 21 +++ .../linux/codex-manifest-linux-arm64.json | 21 +++ 4 files changed, 289 insertions(+) create mode 100644 internal/headless/runtime.go create mode 100644 internal/headless/runtime_test.go create mode 100644 packaging/linux/codex-manifest-linux-amd64.json create mode 100644 packaging/linux/codex-manifest-linux-arm64.json diff --git a/internal/headless/runtime.go b/internal/headless/runtime.go new file mode 100644 index 0000000..d3603cc --- /dev/null +++ b/internal/headless/runtime.go @@ -0,0 +1,86 @@ +package headless + +import ( + "context" + "os/exec" + "path/filepath" + "runtime" + + "github.com/agentserver/agentserver-pkg/internal/codexruntime" + "github.com/agentserver/agentserver-pkg/internal/paths" +) + +type Package struct { + AgentserverExe string + PackageDir string + DriverAgent string + SlaveAgent string +} + +type CodexRuntime struct { + Path string + Source string +} + +type CodexResolveOptions struct { + Paths paths.Paths + Package Package + LookPath func(string) (string, error) + EnsureRuntime func(context.Context, string, string, string) (string, error) +} + +func PackagePaths(agentserverExe string) Package { + packageDir := filepath.Dir(agentserverExe) + return Package{ + AgentserverExe: agentserverExe, + PackageDir: packageDir, + DriverAgent: filepath.Join(packageDir, packageExeName("driver-agent")), + SlaveAgent: filepath.Join(packageDir, packageExeName("slave-agent")), + } +} + +func (p Package) CodexManifestPath() string { + name := "codex-manifest-linux-amd64.json" + if runtime.GOARCH == "arm64" { + name = "codex-manifest-linux-arm64.json" + } + return filepath.Join(p.PackageDir, name) +} + +func ResolveCodex(ctx context.Context, opts CodexResolveOptions) (CodexRuntime, error) { + lookPath := opts.LookPath + if lookPath == nil { + lookPath = exec.LookPath + } + if codexPath, err := lookPath("codex"); err == nil { + return CodexRuntime{Path: codexPath, Source: "path"}, nil + } + + ensureRuntime := opts.EnsureRuntime + if ensureRuntime == nil { + ensureRuntime = func(ctx context.Context, manifestPath, destRoot, cacheDir string) (string, error) { + res, err := codexruntime.Ensure(ctx, codexruntime.Options{ + ManifestPath: manifestPath, + DestRoot: destRoot, + CacheDir: cacheDir, + }) + if err != nil { + return "", err + } + return res.CodexExe, nil + } + } + destRoot := filepath.Dir(filepath.Dir(opts.Paths.CodexExePath)) + codexPath, err := ensureRuntime(ctx, opts.Package.CodexManifestPath(), destRoot, opts.Paths.CacheDir) + if err != nil { + return CodexRuntime{}, err + } + return CodexRuntime{Path: codexPath, Source: "managed"}, nil +} + +func packageExeName(name string) string { + if runtime.GOOS == "windows" { + return name + ".exe" + } + return name +} diff --git a/internal/headless/runtime_test.go b/internal/headless/runtime_test.go new file mode 100644 index 0000000..1189666 --- /dev/null +++ b/internal/headless/runtime_test.go @@ -0,0 +1,161 @@ +package headless + +import ( + "context" + "errors" + "os/exec" + "path/filepath" + "runtime" + "testing" + + "github.com/agentserver/agentserver-pkg/internal/paths" +) + +func TestPackagePathsResolvesSiblingBinaries(t *testing.T) { + exe := filepath.Join(string(filepath.Separator), "opt", "agentserver", exeName("agentserver")) + + pkg := PackagePaths(exe) + + wantDir := filepath.Dir(exe) + if pkg.AgentserverExe != exe { + t.Fatalf("AgentserverExe=%q, want %q", pkg.AgentserverExe, exe) + } + if pkg.PackageDir != wantDir { + t.Fatalf("PackageDir=%q, want %q", pkg.PackageDir, wantDir) + } + if pkg.DriverAgent != filepath.Join(wantDir, exeName("driver-agent")) { + t.Fatalf("DriverAgent=%q", pkg.DriverAgent) + } + if pkg.SlaveAgent != filepath.Join(wantDir, exeName("slave-agent")) { + t.Fatalf("SlaveAgent=%q", pkg.SlaveAgent) + } +} + +func TestResolveCodexPrefersPathCodex(t *testing.T) { + ctx := context.Background() + pathCodex := filepath.Join(t.TempDir(), exeName("codex")) + ensureCalled := false + + got, err := ResolveCodex(ctx, CodexResolveOptions{ + Paths: paths.Paths{ + CodexExePath: filepath.Join(t.TempDir(), "bin-root", "bin", exeName("codex")), + CacheDir: filepath.Join(t.TempDir(), "cache"), + }, + Package: PackagePaths(filepath.Join(string(filepath.Separator), "opt", "agentserver", exeName("agentserver"))), + LookPath: func(name string) (string, error) { + if name != "codex" { + t.Fatalf("LookPath name=%q, want codex", name) + } + return pathCodex, nil + }, + EnsureRuntime: func(context.Context, string, string, string) (string, error) { + ensureCalled = true + return "", nil + }, + }) + if err != nil { + t.Fatal(err) + } + if ensureCalled { + t.Fatal("EnsureRuntime called even though codex was on PATH") + } + if got.Path != pathCodex { + t.Fatalf("Path=%q, want %q", got.Path, pathCodex) + } + if got.Source != "path" { + t.Fatalf("Source=%q, want path", got.Source) + } +} + +func TestResolveCodexEnsuresManagedRuntimeWhenPathMissing(t *testing.T) { + ctx := context.Background() + temp := t.TempDir() + codexExe := filepath.Join(temp, "bin-root", "bin", exeName("codex")) + cacheDir := filepath.Join(temp, "cache") + pkg := PackagePaths(filepath.Join(string(filepath.Separator), "opt", "agentserver", exeName("agentserver"))) + managedCodex := filepath.Join(temp, "managed", exeName("codex")) + var gotManifest, gotDestRoot, gotCacheDir string + + got, err := ResolveCodex(ctx, CodexResolveOptions{ + Paths: paths.Paths{ + CodexExePath: codexExe, + CacheDir: cacheDir, + }, + Package: pkg, + LookPath: func(name string) (string, error) { + if name != "codex" { + t.Fatalf("LookPath name=%q, want codex", name) + } + return "", exec.ErrNotFound + }, + EnsureRuntime: func(_ context.Context, manifestPath, destRoot, cache string) (string, error) { + gotManifest = manifestPath + gotDestRoot = destRoot + gotCacheDir = cache + return managedCodex, nil + }, + }) + if err != nil { + t.Fatal(err) + } + if got.Path != managedCodex { + t.Fatalf("Path=%q, want %q", got.Path, managedCodex) + } + if got.Source != "managed" { + t.Fatalf("Source=%q, want managed", got.Source) + } + if gotManifest != pkg.CodexManifestPath() { + t.Fatalf("manifest=%q, want %q", gotManifest, pkg.CodexManifestPath()) + } + if gotDestRoot != filepath.Dir(filepath.Dir(codexExe)) { + t.Fatalf("destRoot=%q, want %q", gotDestRoot, filepath.Dir(filepath.Dir(codexExe))) + } + if gotCacheDir != cacheDir { + t.Fatalf("cacheDir=%q, want %q", gotCacheDir, cacheDir) + } +} + +func TestLinuxCodexManifestPath(t *testing.T) { + pkg := PackagePaths(filepath.Join(string(filepath.Separator), "opt", "agentserver", exeName("agentserver"))) + + got := pkg.CodexManifestPath() + + var want string + switch runtime.GOARCH { + case "arm64": + want = filepath.Join(string(filepath.Separator), "opt", "agentserver", "codex-manifest-linux-arm64.json") + default: + want = filepath.Join(string(filepath.Separator), "opt", "agentserver", "codex-manifest-linux-amd64.json") + } + if got != want { + t.Fatalf("CodexManifestPath=%q, want %q", got, want) + } +} + +func TestResolveCodexReturnsManagedError(t *testing.T) { + wantErr := errors.New("install failed") + + _, err := ResolveCodex(context.Background(), CodexResolveOptions{ + Paths: paths.Paths{ + CodexExePath: filepath.Join(t.TempDir(), "bin-root", "bin", exeName("codex")), + CacheDir: filepath.Join(t.TempDir(), "cache"), + }, + Package: PackagePaths(filepath.Join(string(filepath.Separator), "opt", "agentserver", exeName("agentserver"))), + LookPath: func(string) (string, error) { + return "", exec.ErrNotFound + }, + EnsureRuntime: func(context.Context, string, string, string) (string, error) { + return "", wantErr + }, + }) + if !errors.Is(err, wantErr) { + t.Fatalf("err=%v, want %v", err, wantErr) + } +} + +func exeName(name string) string { + if runtime.GOOS == "windows" { + return name + ".exe" + } + return name +} diff --git a/packaging/linux/codex-manifest-linux-amd64.json b/packaging/linux/codex-manifest-linux-amd64.json new file mode 100644 index 0000000..7d4784e --- /dev/null +++ b/packaging/linux/codex-manifest-linux-amd64.json @@ -0,0 +1,21 @@ +{ + "package": "@openai/codex", + "platform": "linux-x64", + "pinned_version": "0.139.0-linux-x64", + "strip_prefix": "vendor/x86_64-unknown-linux-musl/", + "codex_exe": "bin/codex", + "required_files": [ + "bin/codex", + "codex-path/rg", + "codex-resources/bwrap", + "codex-resources/zsh/bin/zsh" + ], + "pinned": { + "integrity": "sha512-14UgzDS+X4crkvdt6S02A/ZZOrS8ZyWiuTRpguCtnhNamb7unSuDxy86BWgpAl3sqiTaN2CP8VLyp2ohQ8Nbzw==", + "shasum": "d272f0fa7bd1daf091c419fa22411362cc6b418e", + "urls": [ + "https://registry.npmmirror.com/@openai/codex/-/codex-0.139.0-linux-x64.tgz", + "https://npmreg.proxy.ustclug.org/@openai/codex/-/codex-0.139.0-linux-x64.tgz" + ] + } +} diff --git a/packaging/linux/codex-manifest-linux-arm64.json b/packaging/linux/codex-manifest-linux-arm64.json new file mode 100644 index 0000000..08575c7 --- /dev/null +++ b/packaging/linux/codex-manifest-linux-arm64.json @@ -0,0 +1,21 @@ +{ + "package": "@openai/codex", + "platform": "linux-arm64", + "pinned_version": "0.139.0-linux-arm64", + "strip_prefix": "vendor/aarch64-unknown-linux-musl/", + "codex_exe": "bin/codex", + "required_files": [ + "bin/codex", + "codex-path/rg", + "codex-resources/bwrap", + "codex-resources/zsh/bin/zsh" + ], + "pinned": { + "integrity": "sha512-tBQE5lZciRHeWZGuURgjP9S717MvTIpQMc593+DNxY2LQxozkngOkzFSQd1+/UmQKGrCqdFLu5irIwPXpSZyEw==", + "shasum": "b94e071a7148e21e12d5deea1c6476255f96384b", + "urls": [ + "https://registry.npmmirror.com/@openai/codex/-/codex-0.139.0-linux-arm64.tgz", + "https://npmreg.proxy.ustclug.org/@openai/codex/-/codex-0.139.0-linux-arm64.tgz" + ] + } +} From 03bac6101e2843ac8627788f0b4fc071e4800df7 Mon Sep 17 00:00:00 2001 From: Zishu Yu Date: Sun, 14 Jun 2026 00:17:10 +0800 Subject: [PATCH 14/40] fix(headless): validate managed codex runtime paths --- internal/headless/runtime.go | 17 +++++++ internal/headless/runtime_test.go | 82 +++++++++++++++++++++++++++++++ 2 files changed, 99 insertions(+) diff --git a/internal/headless/runtime.go b/internal/headless/runtime.go index d3603cc..22448fe 100644 --- a/internal/headless/runtime.go +++ b/internal/headless/runtime.go @@ -2,6 +2,7 @@ package headless import ( "context" + "errors" "os/exec" "path/filepath" "runtime" @@ -55,6 +56,9 @@ func ResolveCodex(ctx context.Context, opts CodexResolveOptions) (CodexRuntime, if codexPath, err := lookPath("codex"); err == nil { return CodexRuntime{Path: codexPath, Source: "path"}, nil } + if err := validateManagedCodexOptions(opts); err != nil { + return CodexRuntime{}, err + } ensureRuntime := opts.EnsureRuntime if ensureRuntime == nil { @@ -78,6 +82,19 @@ func ResolveCodex(ctx context.Context, opts CodexResolveOptions) (CodexRuntime, return CodexRuntime{Path: codexPath, Source: "managed"}, nil } +func validateManagedCodexOptions(opts CodexResolveOptions) error { + if opts.Paths.CodexExePath == "" { + return errors.New("CodexExePath required") + } + if opts.Paths.CacheDir == "" { + return errors.New("CacheDir required") + } + if opts.Package.PackageDir == "" { + return errors.New("PackageDir required") + } + return nil +} + func packageExeName(name string) string { if runtime.GOOS == "windows" { return name + ".exe" diff --git a/internal/headless/runtime_test.go b/internal/headless/runtime_test.go index 1189666..88ed99c 100644 --- a/internal/headless/runtime_test.go +++ b/internal/headless/runtime_test.go @@ -6,8 +6,10 @@ import ( "os/exec" "path/filepath" "runtime" + "strings" "testing" + "github.com/agentserver/agentserver-pkg/internal/codexruntime" "github.com/agentserver/agentserver-pkg/internal/paths" ) @@ -115,6 +117,69 @@ func TestResolveCodexEnsuresManagedRuntimeWhenPathMissing(t *testing.T) { } } +func TestResolveCodexRequiresManagedRuntimePathsWhenPathMissing(t *testing.T) { + validPaths := paths.Paths{ + CodexExePath: filepath.Join(t.TempDir(), "bin-root", "bin", exeName("codex")), + CacheDir: filepath.Join(t.TempDir(), "cache"), + } + validPackage := PackagePaths(filepath.Join(string(filepath.Separator), "opt", "agentserver", exeName("agentserver"))) + tests := []struct { + name string + paths paths.Paths + pkg Package + wantErr string + }{ + { + name: "codex exe path", + paths: paths.Paths{ + CacheDir: validPaths.CacheDir, + }, + pkg: validPackage, + wantErr: "CodexExePath required", + }, + { + name: "cache dir", + paths: paths.Paths{ + CodexExePath: validPaths.CodexExePath, + }, + pkg: validPackage, + wantErr: "CacheDir required", + }, + { + name: "package dir", + paths: validPaths, + pkg: Package{}, + wantErr: "PackageDir required", + }, + } + for _, tt := range tests { + t.Run(tt.name, func(t *testing.T) { + ensureCalled := false + + _, err := ResolveCodex(context.Background(), CodexResolveOptions{ + Paths: tt.paths, + Package: tt.pkg, + LookPath: func(string) (string, error) { + return "", exec.ErrNotFound + }, + EnsureRuntime: func(context.Context, string, string, string) (string, error) { + ensureCalled = true + return "", nil + }, + }) + if err == nil { + t.Fatal("expected error") + } + if !strings.Contains(err.Error(), tt.wantErr) { + t.Fatalf("err=%v, want containing %q", err, tt.wantErr) + } + if ensureCalled { + t.Fatal("EnsureRuntime called with invalid managed runtime paths") + } + }) + } +} + func TestLinuxCodexManifestPath(t *testing.T) { pkg := PackagePaths(filepath.Join(string(filepath.Separator), "opt", "agentserver", exeName("agentserver"))) @@ -132,6 +197,23 @@ func TestLinuxCodexManifestPath(t *testing.T) { } } +func TestLinuxCodexManifestsLoad(t *testing.T) { + for _, name := range []string{ + "codex-manifest-linux-amd64.json", + "codex-manifest-linux-arm64.json", + } { + t.Run(name, func(t *testing.T) { + m, err := codexruntime.LoadManifest(filepath.Join("..", "..", "packaging", "linux", name)) + if err != nil { + t.Fatal(err) + } + if m.Package != "@openai/codex" { + t.Fatalf("Package=%q", m.Package) + } + }) + } +} + func TestResolveCodexReturnsManagedError(t *testing.T) { wantErr := errors.New("install failed") From d7676eb1b8b02d05df1144cc55ff8f2ca8714120 Mon Sep 17 00:00:00 2001 From: Zishu Yu Date: Sun, 14 Jun 2026 00:22:43 +0800 Subject: [PATCH 15/40] fix(headless): reject relative managed codex paths --- internal/headless/runtime.go | 9 ++++++ internal/headless/runtime_test.go | 52 +++++++++++++++++++++++++++++++ 2 files changed, 61 insertions(+) diff --git a/internal/headless/runtime.go b/internal/headless/runtime.go index 22448fe..049a8a7 100644 --- a/internal/headless/runtime.go +++ b/internal/headless/runtime.go @@ -86,12 +86,21 @@ func validateManagedCodexOptions(opts CodexResolveOptions) error { if opts.Paths.CodexExePath == "" { return errors.New("CodexExePath required") } + if !filepath.IsAbs(opts.Paths.CodexExePath) { + return errors.New("CodexExePath must be absolute") + } if opts.Paths.CacheDir == "" { return errors.New("CacheDir required") } + if !filepath.IsAbs(opts.Paths.CacheDir) { + return errors.New("CacheDir must be absolute") + } if opts.Package.PackageDir == "" { return errors.New("PackageDir required") } + if !filepath.IsAbs(opts.Package.PackageDir) { + return errors.New("PackageDir must be absolute") + } return nil } diff --git a/internal/headless/runtime_test.go b/internal/headless/runtime_test.go index 88ed99c..4d7b960 100644 --- a/internal/headless/runtime_test.go +++ b/internal/headless/runtime_test.go @@ -151,6 +151,58 @@ func TestResolveCodexRequiresManagedRuntimePathsWhenPathMissing(t *testing.T) { pkg: Package{}, wantErr: "PackageDir required", }, + { + name: "relative codex exe path", + paths: paths.Paths{ + CodexExePath: filepath.Join("bin-root", "bin", exeName("codex")), + CacheDir: validPaths.CacheDir, + }, + pkg: validPackage, + wantErr: "CodexExePath", + }, + { + name: "parent relative codex exe path", + paths: paths.Paths{ + CodexExePath: filepath.Join("..", "bin-root", "bin", exeName("codex")), + CacheDir: validPaths.CacheDir, + }, + pkg: validPackage, + wantErr: "CodexExePath", + }, + { + name: "relative cache dir", + paths: paths.Paths{ + CodexExePath: validPaths.CodexExePath, + CacheDir: "cache", + }, + pkg: validPackage, + wantErr: "CacheDir", + }, + { + name: "parent relative cache dir", + paths: paths.Paths{ + CodexExePath: validPaths.CodexExePath, + CacheDir: filepath.Join("..", "cache"), + }, + pkg: validPackage, + wantErr: "CacheDir", + }, + { + name: "relative package dir", + paths: validPaths, + pkg: Package{ + PackageDir: filepath.Join("packaging", "linux"), + }, + wantErr: "PackageDir", + }, + { + name: "parent relative package dir", + paths: validPaths, + pkg: Package{ + PackageDir: filepath.Join("..", "packaging", "linux"), + }, + wantErr: "PackageDir", + }, } for _, tt := range tests { t.Run(tt.name, func(t *testing.T) { From ee261ff65f71e162d6cc7587ba7babbcc9704632 Mon Sep 17 00:00:00 2001 From: Zishu Yu Date: Sun, 14 Jun 2026 00:32:33 +0800 Subject: [PATCH 16/40] fix(headless): validate codex managed runtime layout --- internal/headless/runtime.go | 62 +++++++++++++-- internal/headless/runtime_test.go | 126 ++++++++++++++++++++++++++++-- 2 files changed, 176 insertions(+), 12 deletions(-) diff --git a/internal/headless/runtime.go b/internal/headless/runtime.go index 049a8a7..9f26427 100644 --- a/internal/headless/runtime.go +++ b/internal/headless/runtime.go @@ -3,6 +3,7 @@ package headless import ( "context" "errors" + "fmt" "os/exec" "path/filepath" "runtime" @@ -41,11 +42,33 @@ func PackagePaths(agentserverExe string) Package { } func (p Package) CodexManifestPath() string { - name := "codex-manifest-linux-amd64.json" - if runtime.GOARCH == "arm64" { - name = "codex-manifest-linux-arm64.json" + path, err := p.codexManifestPath(runtime.GOOS, runtime.GOARCH) + if err != nil { + return "" + } + return path +} + +func (p Package) codexManifestPath(goos, goarch string) (string, error) { + name, err := codexManifestName(goos, goarch) + if err != nil { + return "", err + } + return filepath.Join(p.PackageDir, name), nil +} + +func codexManifestName(goos, goarch string) (string, error) { + if goos != "linux" { + return "", fmt.Errorf("unsupported Codex runtime platform %s/%s", goos, goarch) + } + switch goarch { + case "amd64": + return "codex-manifest-linux-amd64.json", nil + case "arm64": + return "codex-manifest-linux-arm64.json", nil + default: + return "", fmt.Errorf("unsupported Codex runtime platform %s/%s", goos, goarch) } - return filepath.Join(p.PackageDir, name) } func ResolveCodex(ctx context.Context, opts CodexResolveOptions) (CodexRuntime, error) { @@ -74,8 +97,15 @@ func ResolveCodex(ctx context.Context, opts CodexResolveOptions) (CodexRuntime, return res.CodexExe, nil } } - destRoot := filepath.Dir(filepath.Dir(opts.Paths.CodexExePath)) - codexPath, err := ensureRuntime(ctx, opts.Package.CodexManifestPath(), destRoot, opts.Paths.CacheDir) + manifestPath, err := opts.Package.codexManifestPath(runtime.GOOS, runtime.GOARCH) + if err != nil { + return CodexRuntime{}, err + } + destRoot, err := managedCodexDestRoot(opts.Paths.CodexExePath) + if err != nil { + return CodexRuntime{}, err + } + codexPath, err := ensureRuntime(ctx, manifestPath, destRoot, opts.Paths.CacheDir) if err != nil { return CodexRuntime{}, err } @@ -89,6 +119,9 @@ func validateManagedCodexOptions(opts CodexResolveOptions) error { if !filepath.IsAbs(opts.Paths.CodexExePath) { return errors.New("CodexExePath must be absolute") } + if _, err := managedCodexDestRoot(opts.Paths.CodexExePath); err != nil { + return err + } if opts.Paths.CacheDir == "" { return errors.New("CacheDir required") } @@ -104,6 +137,23 @@ func validateManagedCodexOptions(opts CodexResolveOptions) error { return nil } +func managedCodexDestRoot(codexExePath string) (string, error) { + cleanPath := filepath.Clean(codexExePath) + codexExeName := packageExeName("codex") + if filepath.Base(cleanPath) != codexExeName { + return "", fmt.Errorf("CodexExePath must use managed runtime layout /bin/%s", codexExeName) + } + binDir := filepath.Dir(cleanPath) + if filepath.Base(binDir) != "bin" { + return "", fmt.Errorf("CodexExePath must use managed runtime layout /bin/%s", codexExeName) + } + destRoot := filepath.Dir(binDir) + if destRoot == string(filepath.Separator) { + return "", errors.New("CodexExePath managed runtime root must not be filesystem root") + } + return destRoot, nil +} + func packageExeName(name string) string { if runtime.GOOS == "windows" { return name + ".exe" diff --git a/internal/headless/runtime_test.go b/internal/headless/runtime_test.go index 4d7b960..573c276 100644 --- a/internal/headless/runtime_test.go +++ b/internal/headless/runtime_test.go @@ -169,6 +169,42 @@ func TestResolveCodexRequiresManagedRuntimePathsWhenPathMissing(t *testing.T) { pkg: validPackage, wantErr: "CodexExePath", }, + { + name: "codex exe path at filesystem root", + paths: paths.Paths{ + CodexExePath: filepath.Join(string(filepath.Separator), exeName("codex")), + CacheDir: validPaths.CacheDir, + }, + pkg: validPackage, + wantErr: "CodexExePath", + }, + { + name: "codex exe path without bin directory", + paths: paths.Paths{ + CodexExePath: filepath.Join(string(filepath.Separator), "tmp", exeName("codex")), + CacheDir: validPaths.CacheDir, + }, + pkg: validPackage, + wantErr: "CodexExePath", + }, + { + name: "codex exe path with wrong executable name", + paths: paths.Paths{ + CodexExePath: filepath.Join(string(filepath.Separator), "tmp", "bin", "not-codex"), + CacheDir: validPaths.CacheDir, + }, + pkg: validPackage, + wantErr: "CodexExePath", + }, + { + name: "codex exe path derives root dest root", + paths: paths.Paths{ + CodexExePath: filepath.Join(string(filepath.Separator), "bin", exeName("codex")), + CacheDir: validPaths.CacheDir, + }, + pkg: validPackage, + wantErr: "CodexExePath", + }, { name: "relative cache dir", paths: paths.Paths{ @@ -232,6 +268,59 @@ func TestResolveCodexRequiresManagedRuntimePathsWhenPathMissing(t *testing.T) { } } +func TestCodexManifestNameSupportsOnlyLinuxCodexRuntimes(t *testing.T) { + tests := []struct { + goos string + goarch string + wantName string + wantErr string + }{ + { + goos: "linux", + goarch: "amd64", + wantName: "codex-manifest-linux-amd64.json", + }, + { + goos: "linux", + goarch: "arm64", + wantName: "codex-manifest-linux-arm64.json", + }, + { + goos: "linux", + goarch: "riscv64", + wantErr: "unsupported Codex runtime platform linux/riscv64", + }, + { + goos: "darwin", + goarch: "amd64", + wantErr: "unsupported Codex runtime platform darwin/amd64", + }, + } + for _, tt := range tests { + t.Run(tt.goos+"/"+tt.goarch, func(t *testing.T) { + got, err := codexManifestName(tt.goos, tt.goarch) + if tt.wantErr != "" { + if err == nil { + t.Fatal("expected error") + } + if !strings.Contains(err.Error(), tt.wantErr) { + t.Fatalf("err=%v, want containing %q", err, tt.wantErr) + } + if got != "" { + t.Fatalf("name=%q, want empty", got) + } + return + } + if err != nil { + t.Fatal(err) + } + if got != tt.wantName { + t.Fatalf("name=%q, want %q", got, tt.wantName) + } + }) + } +} + func TestLinuxCodexManifestPath(t *testing.T) { pkg := PackagePaths(filepath.Join(string(filepath.Separator), "opt", "agentserver", exeName("agentserver"))) @@ -250,18 +339,43 @@ func TestLinuxCodexManifestPath(t *testing.T) { } func TestLinuxCodexManifestsLoad(t *testing.T) { - for _, name := range []string{ - "codex-manifest-linux-amd64.json", - "codex-manifest-linux-arm64.json", - } { - t.Run(name, func(t *testing.T) { - m, err := codexruntime.LoadManifest(filepath.Join("..", "..", "packaging", "linux", name)) + tests := []struct { + name string + platform string + pinnedVersion string + stripPrefix string + }{ + { + name: "codex-manifest-linux-amd64.json", + platform: "linux-x64", + pinnedVersion: "0.139.0-linux-x64", + stripPrefix: "vendor/x86_64-unknown-linux-musl/", + }, + { + name: "codex-manifest-linux-arm64.json", + platform: "linux-arm64", + pinnedVersion: "0.139.0-linux-arm64", + stripPrefix: "vendor/aarch64-unknown-linux-musl/", + }, + } + for _, tt := range tests { + t.Run(tt.name, func(t *testing.T) { + m, err := codexruntime.LoadManifest(filepath.Join("..", "..", "packaging", "linux", tt.name)) if err != nil { t.Fatal(err) } if m.Package != "@openai/codex" { t.Fatalf("Package=%q", m.Package) } + if m.Platform != tt.platform { + t.Fatalf("Platform=%q, want %q", m.Platform, tt.platform) + } + if m.PinnedVersion != tt.pinnedVersion { + t.Fatalf("PinnedVersion=%q, want %q", m.PinnedVersion, tt.pinnedVersion) + } + if m.StripPrefix != tt.stripPrefix { + t.Fatalf("StripPrefix=%q, want %q", m.StripPrefix, tt.stripPrefix) + } }) } } From bc47efe78bc8e4436699e7daa5498c659afd0184 Mon Sep 17 00:00:00 2001 From: Zishu Yu Date: Sun, 14 Jun 2026 00:40:16 +0800 Subject: [PATCH 17/40] test(headless): gate linux codex runtime tests --- internal/headless/runtime_test.go | 45 ++++++++++++++++++++++++------- 1 file changed, 35 insertions(+), 10 deletions(-) diff --git a/internal/headless/runtime_test.go b/internal/headless/runtime_test.go index 573c276..85d74df 100644 --- a/internal/headless/runtime_test.go +++ b/internal/headless/runtime_test.go @@ -70,6 +70,8 @@ func TestResolveCodexPrefersPathCodex(t *testing.T) { } func TestResolveCodexEnsuresManagedRuntimeWhenPathMissing(t *testing.T) { + skipUnsupportedCodexRuntime(t) + ctx := context.Background() temp := t.TempDir() codexExe := filepath.Join(temp, "bin-root", "bin", exeName("codex")) @@ -324,17 +326,29 @@ func TestCodexManifestNameSupportsOnlyLinuxCodexRuntimes(t *testing.T) { func TestLinuxCodexManifestPath(t *testing.T) { pkg := PackagePaths(filepath.Join(string(filepath.Separator), "opt", "agentserver", exeName("agentserver"))) - got := pkg.CodexManifestPath() - - var want string - switch runtime.GOARCH { - case "arm64": - want = filepath.Join(string(filepath.Separator), "opt", "agentserver", "codex-manifest-linux-arm64.json") - default: - want = filepath.Join(string(filepath.Separator), "opt", "agentserver", "codex-manifest-linux-amd64.json") + tests := []struct { + goarch string + want string + }{ + { + goarch: "amd64", + want: filepath.Join(string(filepath.Separator), "opt", "agentserver", "codex-manifest-linux-amd64.json"), + }, + { + goarch: "arm64", + want: filepath.Join(string(filepath.Separator), "opt", "agentserver", "codex-manifest-linux-arm64.json"), + }, } - if got != want { - t.Fatalf("CodexManifestPath=%q, want %q", got, want) + for _, tt := range tests { + t.Run("linux/"+tt.goarch, func(t *testing.T) { + got, err := pkg.codexManifestPath("linux", tt.goarch) + if err != nil { + t.Fatal(err) + } + if got != tt.want { + t.Fatalf("CodexManifestPath=%q, want %q", got, tt.want) + } + }) } } @@ -381,6 +395,8 @@ func TestLinuxCodexManifestsLoad(t *testing.T) { } func TestResolveCodexReturnsManagedError(t *testing.T) { + skipUnsupportedCodexRuntime(t) + wantErr := errors.New("install failed") _, err := ResolveCodex(context.Background(), CodexResolveOptions{ @@ -401,6 +417,15 @@ func TestResolveCodexReturnsManagedError(t *testing.T) { } } +func skipUnsupportedCodexRuntime(t *testing.T) { + t.Helper() + + if runtime.GOOS == "linux" && (runtime.GOARCH == "amd64" || runtime.GOARCH == "arm64") { + return + } + t.Skipf("managed Codex runtime is unsupported on %s/%s", runtime.GOOS, runtime.GOARCH) +} + func exeName(name string) string { if runtime.GOOS == "windows" { return name + ".exe" From 4be3857205819c9b27a033d7faa78022f3dfedfc Mon Sep 17 00:00:00 2001 From: Zishu Yu Date: Sun, 14 Jun 2026 00:51:06 +0800 Subject: [PATCH 18/40] feat(headless): reuse slaves by canonical folder --- internal/slave/registry.go | 150 +++++++++++++++++++++++++------- internal/slave/registry_test.go | 100 +++++++++++++++++++++ 2 files changed, 217 insertions(+), 33 deletions(-) diff --git a/internal/slave/registry.go b/internal/slave/registry.go index 24c287b..db5a0b5 100644 --- a/internal/slave/registry.go +++ b/internal/slave/registry.go @@ -46,6 +46,28 @@ type CreateInput struct { Name string } +func CanonicalFolder(folder string) (string, error) { + folder = strings.TrimSpace(folder) + if folder == "" { + return "", fmt.Errorf("%w: folder required", ErrInvalidCreateInput) + } + abs, err := filepath.Abs(folder) + if err != nil { + return "", fmt.Errorf("%w: folder required: %v", ErrInvalidCreateInput, err) + } + info, err := os.Stat(abs) + if err != nil { + return "", fmt.Errorf("%w: folder unavailable: %w", ErrInvalidCreateInput, err) + } + if !info.IsDir() { + return "", fmt.Errorf("%w: folder is not a directory: %s", ErrInvalidCreateInput, abs) + } + if realPath, err := filepath.EvalSymlinks(abs); err == nil { + abs = realPath + } + return filepath.Clean(abs), nil +} + var ( ErrInvalidCreateInput = errors.New("invalid slave create input") ErrSlaveConflict = errors.New("slave conflict") @@ -120,20 +142,9 @@ func (r *Registry) Create(machine Machine, in CreateInput) (Slave, error) { if strings.TrimSpace(machine.MachineID) == "" || strings.TrimSpace(machine.ComputerName) == "" { return Slave{}, fmt.Errorf("machine identity required") } - folderInput := strings.TrimSpace(in.Folder) - if folderInput == "" { - return Slave{}, fmt.Errorf("%w: folder required", ErrInvalidCreateInput) - } - folder, err := filepath.Abs(folderInput) - if err != nil { - return Slave{}, fmt.Errorf("%w: folder required: %v", ErrInvalidCreateInput, err) - } - info, err := os.Stat(folder) + folder, err := CanonicalFolder(in.Folder) if err != nil { - return Slave{}, fmt.Errorf("%w: folder unavailable: %w", ErrInvalidCreateInput, err) - } - if !info.IsDir() { - return Slave{}, fmt.Errorf("%w: folder is not a directory: %s", ErrInvalidCreateInput, folder) + return Slave{}, err } name := strings.TrimSpace(in.Name) if name == "" { @@ -143,10 +154,6 @@ func (r *Registry) Create(machine Machine, in CreateInput) (Slave, error) { return Slave{}, fmt.Errorf("%w: %w", ErrInvalidCreateInput, err) } displayName := machine.ComputerName + "-" + name - id, err := uuid.NewRandom() - if err != nil { - return Slave{}, fmt.Errorf("generate slave id: %w", err) - } r.mu.Lock() defer r.mu.Unlock() @@ -159,24 +166,76 @@ func (r *Registry) Create(machine Machine, in CreateInput) (Slave, error) { return Slave{}, fmt.Errorf("%w: slave display name already exists: %s", ErrSlaveConflict, displayName) } } - now := time.Now().UTC() - dir := filepath.Join(r.slavesDir, id.String()) - sl := Slave{ - ID: id.String(), - Name: name, - DisplayName: displayName, - Folder: folder, - ConfigPath: filepath.Join(dir, "config.yaml"), - LogPath: filepath.Join(dir, "logs", "slave.log"), - Status: StatusStopped, - CreatedAt: now, - UpdatedAt: now, + return r.createLocked(all, name, displayName, folder) +} + +func (r *Registry) FindByFolder(folder string) (Slave, bool, error) { + canonical, err := CanonicalFolder(folder) + if err != nil { + return Slave{}, false, err } - all = append(all, sl) - if err := r.saveLocked(all); err != nil { - return Slave{}, err + r.mu.Lock() + defer r.mu.Unlock() + + all, err := r.loadLocked() + if err != nil { + return Slave{}, false, err } - return sl, nil + for _, sl := range all { + existing, err := CanonicalFolder(sl.Folder) + if err != nil { + existing = filepath.Clean(sl.Folder) + } + if existing == canonical { + return sl, true, nil + } + } + return Slave{}, false, nil +} + +func (r *Registry) EnsureForFolder(machine Machine, in CreateInput) (Slave, bool, error) { + folder, err := CanonicalFolder(in.Folder) + if err != nil { + return Slave{}, false, err + } + if strings.TrimSpace(machine.MachineID) == "" || strings.TrimSpace(machine.ComputerName) == "" { + return Slave{}, false, fmt.Errorf("machine identity required") + } + name := strings.TrimSpace(in.Name) + if name == "" { + name = filepath.Base(folder) + } + if err := validateSlaveName(name); err != nil { + return Slave{}, false, fmt.Errorf("%w: %w", ErrInvalidCreateInput, err) + } + displayName := machine.ComputerName + "-" + name + + r.mu.Lock() + defer r.mu.Unlock() + + all, err := r.loadLocked() + if err != nil { + return Slave{}, false, err + } + for _, existing := range all { + existingFolder, err := CanonicalFolder(existing.Folder) + if err != nil { + existingFolder = filepath.Clean(existing.Folder) + } + if existingFolder == folder { + return existing, false, nil + } + } + for _, existing := range all { + if existing.DisplayName == displayName { + return Slave{}, false, fmt.Errorf("%w: slave display name already exists: %s", ErrSlaveConflict, displayName) + } + } + sl, err := r.createLocked(all, name, displayName, folder) + if err != nil { + return Slave{}, false, err + } + return sl, true, nil } func (r *Registry) Update(id string, fn func(*Slave) error) (Slave, error) { @@ -224,6 +283,31 @@ func (r *Registry) Delete(id string) (Slave, error) { return Slave{}, os.ErrNotExist } +func (r *Registry) createLocked(all []Slave, name, displayName, folder string) (Slave, error) { + id, err := uuid.NewRandom() + if err != nil { + return Slave{}, fmt.Errorf("generate slave id: %w", err) + } + now := time.Now().UTC() + dir := filepath.Join(r.slavesDir, id.String()) + sl := Slave{ + ID: id.String(), + Name: name, + DisplayName: displayName, + Folder: folder, + ConfigPath: filepath.Join(dir, "config.yaml"), + LogPath: filepath.Join(dir, "logs", "slave.log"), + Status: StatusStopped, + CreatedAt: now, + UpdatedAt: now, + } + all = append(all, sl) + if err := r.saveLocked(all); err != nil { + return Slave{}, err + } + return sl, nil +} + func validateSlaveName(name string) error { if strings.TrimSpace(name) == "" { return fmt.Errorf("slave name required") diff --git a/internal/slave/registry_test.go b/internal/slave/registry_test.go index 1d17e3b..9e50664 100644 --- a/internal/slave/registry_test.go +++ b/internal/slave/registry_test.go @@ -108,6 +108,106 @@ func TestRegistryRejectsDuplicateDisplayName(t *testing.T) { } } +func TestRegistryFindByFolderUsesCanonicalPath(t *testing.T) { + dir := t.TempDir() + folder := filepath.Join(dir, "repo") + if err := mkdir(folder); err != nil { + t.Fatal(err) + } + link := filepath.Join(dir, "repo-link") + if err := os.Symlink(folder, link); err != nil { + t.Skipf("symlink unavailable: %v", err) + } + reg := NewRegistry(filepath.Join(dir, "slaves.json"), filepath.Join(dir, "slaves")) + m := Machine{MachineID: "machine-1", ComputerName: "host"} + created, err := reg.Create(m, CreateInput{Folder: folder, Name: "repo"}) + if err != nil { + t.Fatalf("Create: %v", err) + } + + got, ok, err := reg.FindByFolder(link) + if err != nil { + t.Fatalf("FindByFolder: %v", err) + } + if !ok { + t.Fatal("FindByFolder did not find existing slave") + } + if got.ID != created.ID { + t.Fatalf("got ID=%q want %q", got.ID, created.ID) + } +} + +func TestRegistryEnsureForFolderReusesExistingSlave(t *testing.T) { + dir := t.TempDir() + folder := filepath.Join(dir, "repo") + if err := mkdir(folder); err != nil { + t.Fatal(err) + } + reg := NewRegistry(filepath.Join(dir, "slaves.json"), filepath.Join(dir, "slaves")) + m := Machine{MachineID: "machine-1", ComputerName: "host"} + + first, created, err := reg.EnsureForFolder(m, CreateInput{Folder: folder, Name: "repo"}) + if err != nil { + t.Fatalf("EnsureForFolder first: %v", err) + } + if !created { + t.Fatal("first EnsureForFolder should create") + } + second, created, err := reg.EnsureForFolder(m, CreateInput{Folder: folder, Name: "different"}) + if err != nil { + t.Fatalf("EnsureForFolder second: %v", err) + } + if created { + t.Fatal("second EnsureForFolder should reuse") + } + if second.ID != first.ID || second.Name != "repo" { + t.Fatalf("second=%+v first=%+v", second, first) + } +} + +func TestRegistryCreateStoresCanonicalFolderWhenInputIsSymlink(t *testing.T) { + dir := t.TempDir() + folder := filepath.Join(dir, "repo") + if err := mkdir(folder); err != nil { + t.Fatal(err) + } + link := filepath.Join(dir, "repo-link") + if err := os.Symlink(folder, link); err != nil { + t.Skipf("symlink unavailable: %v", err) + } + reg := NewRegistry(filepath.Join(dir, "slaves.json"), filepath.Join(dir, "slaves")) + m := Machine{MachineID: "machine-1", ComputerName: "host"} + + got, err := reg.Create(m, CreateInput{Folder: link, Name: "repo"}) + if err != nil { + t.Fatalf("Create: %v", err) + } + if got.Folder != folder { + t.Fatalf("Folder=%q want %q", got.Folder, folder) + } +} + +func TestRegistryEnsureForFolderRejectsDuplicateDisplayNameForDifferentFolders(t *testing.T) { + dir := t.TempDir() + folderA := filepath.Join(dir, "a") + folderB := filepath.Join(dir, "b") + if err := mkdir(folderA); err != nil { + t.Fatal(err) + } + if err := mkdir(folderB); err != nil { + t.Fatal(err) + } + reg := NewRegistry(filepath.Join(dir, "slaves.json"), filepath.Join(dir, "slaves")) + m := Machine{MachineID: "machine-1", ComputerName: "host"} + + if _, created, err := reg.EnsureForFolder(m, CreateInput{Folder: folderA, Name: "worker"}); err != nil || !created { + t.Fatalf("EnsureForFolder first created=%v err=%v", created, err) + } + if _, _, err := reg.EnsureForFolder(m, CreateInput{Folder: folderB, Name: "worker"}); !errors.Is(err, ErrSlaveConflict) { + t.Fatalf("expected duplicate display name error, got %v", err) + } +} + func TestRegistryConcurrentCreatesPreserveAllSlaves(t *testing.T) { prevProcs := runtime.GOMAXPROCS(0) if prevProcs < 8 { From 267fbb1af3e2cf035733521ea866bf8d084d61cf Mon Sep 17 00:00:00 2001 From: Zishu Yu Date: Sun, 14 Jun 2026 01:00:43 +0800 Subject: [PATCH 19/40] fix(headless): harden canonical slave registry lookup --- internal/slave/registry.go | 80 ++++++++---- internal/slave/registry_test.go | 223 +++++++++++++++++++++++++++++++- 2 files changed, 278 insertions(+), 25 deletions(-) diff --git a/internal/slave/registry.go b/internal/slave/registry.go index db5a0b5..126894d 100644 --- a/internal/slave/registry.go +++ b/internal/slave/registry.go @@ -13,7 +13,10 @@ import ( "github.com/google/uuid" ) -var osMkdirAll = os.MkdirAll +var ( + osMkdirAll = os.MkdirAll + evalSymlinks = filepath.EvalSymlinks +) type Status string @@ -62,12 +65,38 @@ func CanonicalFolder(folder string) (string, error) { if !info.IsDir() { return "", fmt.Errorf("%w: folder is not a directory: %s", ErrInvalidCreateInput, abs) } - if realPath, err := filepath.EvalSymlinks(abs); err == nil { - abs = realPath + realPath, err := evalSymlinks(abs) + if err != nil { + return "", fmt.Errorf("%w: resolve folder symlinks: %w", ErrInvalidCreateInput, err) } + abs = realPath return filepath.Clean(abs), nil } +func defaultFolderName(folder string) (string, error) { + folder = strings.TrimSpace(folder) + abs, err := filepath.Abs(folder) + if err != nil { + return "", fmt.Errorf("%w: folder required: %v", ErrInvalidCreateInput, err) + } + return filepath.Base(filepath.Clean(abs)), nil +} + +func storedFolderKey(folder string) string { + folder = strings.TrimSpace(folder) + if folder == "" { + return "" + } + abs, err := filepath.Abs(folder) + if err != nil { + return filepath.Clean(folder) + } + if realPath, err := evalSymlinks(abs); err == nil { + abs = realPath + } + return filepath.Clean(abs) +} + var ( ErrInvalidCreateInput = errors.New("invalid slave create input") ErrSlaveConflict = errors.New("slave conflict") @@ -148,7 +177,10 @@ func (r *Registry) Create(machine Machine, in CreateInput) (Slave, error) { } name := strings.TrimSpace(in.Name) if name == "" { - name = filepath.Base(folder) + name, err = defaultFolderName(in.Folder) + if err != nil { + return Slave{}, err + } } if err := validateSlaveName(name); err != nil { return Slave{}, fmt.Errorf("%w: %w", ErrInvalidCreateInput, err) @@ -174,19 +206,12 @@ func (r *Registry) FindByFolder(folder string) (Slave, bool, error) { if err != nil { return Slave{}, false, err } - r.mu.Lock() - defer r.mu.Unlock() - - all, err := r.loadLocked() + all, err := r.List() if err != nil { return Slave{}, false, err } for _, sl := range all { - existing, err := CanonicalFolder(sl.Folder) - if err != nil { - existing = filepath.Clean(sl.Folder) - } - if existing == canonical { + if storedFolderKey(sl.Folder) == canonical { return sl, true, nil } } @@ -203,35 +228,44 @@ func (r *Registry) EnsureForFolder(machine Machine, in CreateInput) (Slave, bool } name := strings.TrimSpace(in.Name) if name == "" { - name = filepath.Base(folder) + name, err = defaultFolderName(in.Folder) + if err != nil { + return Slave{}, false, err + } } if err := validateSlaveName(name); err != nil { return Slave{}, false, fmt.Errorf("%w: %w", ErrInvalidCreateInput, err) } displayName := machine.ComputerName + "-" + name + all, err := r.List() + if err != nil { + return Slave{}, false, err + } + for _, existing := range all { + if storedFolderKey(existing.Folder) == folder { + return existing, false, nil + } + } + r.mu.Lock() defer r.mu.Unlock() - all, err := r.loadLocked() + latest, err := r.loadLocked() if err != nil { return Slave{}, false, err } - for _, existing := range all { - existingFolder, err := CanonicalFolder(existing.Folder) - if err != nil { - existingFolder = filepath.Clean(existing.Folder) - } - if existingFolder == folder { + for _, existing := range latest { + if filepath.Clean(existing.Folder) == folder { return existing, false, nil } } - for _, existing := range all { + for _, existing := range latest { if existing.DisplayName == displayName { return Slave{}, false, fmt.Errorf("%w: slave display name already exists: %s", ErrSlaveConflict, displayName) } } - sl, err := r.createLocked(all, name, displayName, folder) + sl, err := r.createLocked(latest, name, displayName, folder) if err != nil { return Slave{}, false, err } diff --git a/internal/slave/registry_test.go b/internal/slave/registry_test.go index 9e50664..a17d987 100644 --- a/internal/slave/registry_test.go +++ b/internal/slave/registry_test.go @@ -1,6 +1,7 @@ package slave import ( + "encoding/json" "errors" "fmt" "os" @@ -137,6 +138,51 @@ func TestRegistryFindByFolderUsesCanonicalPath(t *testing.T) { } } +func TestRegistryFindByFolderSkipsStaleStoredFolderAndFindsLaterMatch(t *testing.T) { + dir := t.TempDir() + folder := filepath.Join(dir, "repo") + if err := mkdir(folder); err != nil { + t.Fatal(err) + } + stale := filepath.Join(dir, "missing") + registryPath := filepath.Join(dir, "slaves.json") + all := []Slave{ + { + ID: "stale", + Name: "stale", + DisplayName: "host-stale", + Folder: stale, + Status: StatusStopped, + }, + { + ID: "valid", + Name: "valid", + DisplayName: "host-valid", + Folder: folder, + Status: StatusStopped, + }, + } + b, err := json.Marshal(all) + if err != nil { + t.Fatal(err) + } + if err := os.WriteFile(registryPath, b, 0o600); err != nil { + t.Fatal(err) + } + reg := NewRegistry(registryPath, filepath.Join(dir, "slaves")) + + got, ok, err := reg.FindByFolder(folder) + if err != nil { + t.Fatalf("FindByFolder: %v", err) + } + if !ok { + t.Fatal("FindByFolder did not find existing slave") + } + if got.ID != "valid" { + t.Fatalf("got ID=%q want valid", got.ID) + } +} + func TestRegistryEnsureForFolderReusesExistingSlave(t *testing.T) { dir := t.TempDir() folder := filepath.Join(dir, "repo") @@ -182,8 +228,114 @@ func TestRegistryCreateStoresCanonicalFolderWhenInputIsSymlink(t *testing.T) { if err != nil { t.Fatalf("Create: %v", err) } - if got.Folder != folder { - t.Fatalf("Folder=%q want %q", got.Folder, folder) + want, err := CanonicalFolder(folder) + if err != nil { + t.Fatalf("CanonicalFolder: %v", err) + } + if got.Folder != want { + t.Fatalf("Folder=%q want %q", got.Folder, want) + } +} + +func TestRegistryCreateDefaultNameUsesSymlinkInputBasename(t *testing.T) { + dir := t.TempDir() + folder := filepath.Join(dir, "repo") + if err := mkdir(folder); err != nil { + t.Fatal(err) + } + link := filepath.Join(dir, "repo-link") + if err := os.Symlink(folder, link); err != nil { + t.Skipf("symlink unavailable: %v", err) + } + reg := NewRegistry(filepath.Join(dir, "slaves.json"), filepath.Join(dir, "slaves")) + m := Machine{MachineID: "machine-1", ComputerName: "host"} + + got, err := reg.Create(m, CreateInput{Folder: link}) + if err != nil { + t.Fatalf("Create: %v", err) + } + if got.Name != "repo-link" { + t.Fatalf("Name=%q want repo-link", got.Name) + } + if got.DisplayName != "host-repo-link" { + t.Fatalf("DisplayName=%q want host-repo-link", got.DisplayName) + } + want, err := CanonicalFolder(folder) + if err != nil { + t.Fatalf("CanonicalFolder: %v", err) + } + if got.Folder != want { + t.Fatalf("Folder=%q want %q", got.Folder, want) + } +} + +func TestRegistryEnsureForFolderDefaultNameUsesSymlinkInputBasenameAndReusesCanonicalTarget(t *testing.T) { + dir := t.TempDir() + folder := filepath.Join(dir, "repo") + if err := mkdir(folder); err != nil { + t.Fatal(err) + } + link := filepath.Join(dir, "repo-link") + if err := os.Symlink(folder, link); err != nil { + t.Skipf("symlink unavailable: %v", err) + } + reg := NewRegistry(filepath.Join(dir, "slaves.json"), filepath.Join(dir, "slaves")) + m := Machine{MachineID: "machine-1", ComputerName: "host"} + + first, created, err := reg.EnsureForFolder(m, CreateInput{Folder: link}) + if err != nil { + t.Fatalf("EnsureForFolder first: %v", err) + } + if !created { + t.Fatal("first EnsureForFolder should create") + } + if first.Name != "repo-link" { + t.Fatalf("first Name=%q want repo-link", first.Name) + } + want, err := CanonicalFolder(folder) + if err != nil { + t.Fatalf("CanonicalFolder: %v", err) + } + if first.Folder != want { + t.Fatalf("first Folder=%q want %q", first.Folder, want) + } + + second, created, err := reg.EnsureForFolder(m, CreateInput{Folder: folder}) + if err != nil { + t.Fatalf("EnsureForFolder second: %v", err) + } + if created { + t.Fatal("second EnsureForFolder should reuse") + } + if second.ID != first.ID { + t.Fatalf("second ID=%q want %q", second.ID, first.ID) + } + if second.Name != "repo-link" { + t.Fatalf("second Name=%q want repo-link", second.Name) + } +} + +func TestCanonicalFolderReturnsInvalidCreateInputWhenEvalSymlinksFails(t *testing.T) { + dir := t.TempDir() + folder := filepath.Join(dir, "repo") + if err := mkdir(folder); err != nil { + t.Fatal(err) + } + wantErr := errors.New("resolve failed") + prev := evalSymlinks + evalSymlinks = func(string) (string, error) { + return "", wantErr + } + t.Cleanup(func() { + evalSymlinks = prev + }) + + _, err := CanonicalFolder(folder) + if !errors.Is(err, ErrInvalidCreateInput) { + t.Fatalf("error=%v, want ErrInvalidCreateInput", err) + } + if !errors.Is(err, wantErr) { + t.Fatalf("error=%v, want wrapped eval error", err) } } @@ -208,6 +360,73 @@ func TestRegistryEnsureForFolderRejectsDuplicateDisplayNameForDifferentFolders(t } } +func TestRegistryConcurrentEnsureForFolderCreatesOneSlave(t *testing.T) { + prevProcs := runtime.GOMAXPROCS(0) + if prevProcs < 8 { + runtime.GOMAXPROCS(8) + t.Cleanup(func() { + runtime.GOMAXPROCS(prevProcs) + }) + } + + dir := t.TempDir() + folder := filepath.Join(dir, "repo") + if err := mkdir(folder); err != nil { + t.Fatal(err) + } + reg := NewRegistry(filepath.Join(dir, "slaves.json"), filepath.Join(dir, "slaves")) + m := Machine{MachineID: "machine-1", ComputerName: "host"} + const workers = 64 + start := make(chan struct{}) + results := make([]Slave, workers) + created := make([]bool, workers) + errs := make([]error, workers) + + var wg sync.WaitGroup + wg.Add(workers) + for i := 0; i < workers; i++ { + i := i + go func() { + defer wg.Done() + <-start + results[i], created[i], errs[i] = reg.EnsureForFolder(m, CreateInput{Folder: folder}) + }() + } + + close(start) + wg.Wait() + + var id string + createCount := 0 + for i, err := range errs { + if err != nil { + t.Fatalf("worker %d EnsureForFolder: %v", i, err) + } + if results[i].ID == "" { + t.Fatalf("worker %d empty ID: %+v", i, results[i]) + } + if id == "" { + id = results[i].ID + } + if results[i].ID != id { + t.Fatalf("worker %d ID=%q want %q", i, results[i].ID, id) + } + if created[i] { + createCount++ + } + } + if createCount != 1 { + t.Fatalf("created count=%d want 1", createCount) + } + all, err := reg.List() + if err != nil { + t.Fatalf("List: %v", err) + } + if len(all) != 1 { + t.Fatalf("registry has %d slaves, want 1: %+v", len(all), all) + } +} + func TestRegistryConcurrentCreatesPreserveAllSlaves(t *testing.T) { prevProcs := runtime.GOMAXPROCS(0) if prevProcs < 8 { From a7e2bf76d2ee4b6e2b5411f78a520d552bb99034 Mon Sep 17 00:00:00 2001 From: Zishu Yu Date: Sun, 14 Jun 2026 01:10:47 +0800 Subject: [PATCH 20/40] feat(headless): run directory slaves in foreground --- internal/headless/slave.go | 182 ++++++++++++++++++++++++++ internal/headless/slave_test.go | 223 ++++++++++++++++++++++++++++++++ 2 files changed, 405 insertions(+) create mode 100644 internal/headless/slave.go create mode 100644 internal/headless/slave_test.go diff --git a/internal/headless/slave.go b/internal/headless/slave.go new file mode 100644 index 0000000..138baa6 --- /dev/null +++ b/internal/headless/slave.go @@ -0,0 +1,182 @@ +package headless + +import ( + "bufio" + "context" + "fmt" + "io" + "os" + "os/exec" + "path/filepath" + "regexp" + "strings" + "sync" + + "github.com/agentserver/agentserver-pkg/internal/paths" + "github.com/agentserver/agentserver-pkg/internal/slave" + "github.com/agentserver/agentserver-pkg/internal/terminalauth" +) + +type NamePrompt func(defaultName string) (string, error) + +type SlaveOptions struct { + Paths paths.Paths + Package Package + WorkDir string + ComputerName string + NamePrompt NamePrompt + Runner SlaveProcessRunner + Stdout io.Writer + QR terminalauth.QRWriter + CodexBin string +} + +type SlaveProcessRequest struct { + Exe string + ConfigPath string + WorkDir string + AuthURL func(string) +} + +type SlaveProcessRunner interface { + Run(context.Context, SlaveProcessRequest) error +} + +func RunSlave(ctx context.Context, opts SlaveOptions) error { + stdout := opts.Stdout + if stdout == nil { + stdout = os.Stdout + } + runner := opts.Runner + if runner == nil { + runner = execSlaveRunner{stdout: stdout} + } + computerName := strings.TrimSpace(opts.ComputerName) + if computerName == "" { + hostname, err := os.Hostname() + if err != nil { + return fmt.Errorf("hostname: %w", err) + } + computerName = hostname + } + + canonicalFolder, err := slave.CanonicalFolder(opts.WorkDir) + if err != nil { + return err + } + machine, err := slave.NewMachineStore(opts.Paths.MachineFile).Ensure(computerName) + if err != nil { + return err + } + registry := slave.NewRegistry(opts.Paths.SlavesFile, opts.Paths.SlavesDir) + + name := "" + if existing, ok, err := registry.FindByFolder(canonicalFolder); err != nil { + return err + } else if ok { + name = existing.Name + } else { + name = filepath.Base(canonicalFolder) + if opts.NamePrompt != nil { + prompted, err := opts.NamePrompt(name) + if err != nil { + return err + } + if strings.TrimSpace(prompted) != "" { + name = prompted + } + } + } + + sl, _, err := registry.EnsureForFolder(machine, slave.CreateInput{Folder: canonicalFolder, Name: name}) + if err != nil { + return err + } + if err := slave.WriteConfig(sl, machine, slave.ConfigInput{CodexBin: opts.CodexBin}); err != nil { + return err + } + + return runner.Run(ctx, SlaveProcessRequest{ + Exe: opts.Package.SlaveAgent, + ConfigPath: sl.ConfigPath, + WorkDir: filepath.Dir(sl.ConfigPath), + AuthURL: func(url string) { + terminalauth.PrintURL(stdout, "Agentserver 登录", url, "", opts.QR) + }, + }) +} + +type execSlaveRunner struct { + stdout io.Writer +} + +var foregroundAuthURLPattern = regexp.MustCompile(`(?i)https?://\S*(?:device|user[_-]code|verification)\S*`) + +func (r execSlaveRunner) Run(ctx context.Context, req SlaveProcessRequest) error { + cmd := exec.CommandContext(ctx, req.Exe, req.ConfigPath) + cmd.Dir = req.WorkDir + + stdout, err := cmd.StdoutPipe() + if err != nil { + return fmt.Errorf("slave stdout pipe: %w", err) + } + stderr, err := cmd.StderrPipe() + if err != nil { + return fmt.Errorf("slave stderr pipe: %w", err) + } + if err := cmd.Start(); err != nil { + return fmt.Errorf("start slave-agent: %w", err) + } + + writer := r.stdout + if writer == nil { + writer = os.Stdout + } + writer = &foregroundLockedWriter{w: writer} + var wg sync.WaitGroup + wg.Add(2) + go func() { + defer wg.Done() + copyForegroundSlaveOutput(stdout, writer, req.AuthURL) + }() + go func() { + defer wg.Done() + copyForegroundSlaveOutput(stderr, writer, req.AuthURL) + }() + + err = cmd.Wait() + wg.Wait() + if ctx.Err() != nil { + return ctx.Err() + } + if err != nil { + return fmt.Errorf("slave-agent exited: %w", err) + } + return nil +} + +func copyForegroundSlaveOutput(r io.Reader, w io.Writer, authURL func(string)) { + scanner := bufio.NewScanner(r) + scanner.Buffer(make([]byte, 1<<16), 1<<20) + for scanner.Scan() { + line := scanner.Text() + _, _ = fmt.Fprintln(w, line) + if authURL != nil && foregroundAuthURLPattern.MatchString(line) { + authURL(foregroundAuthURLPattern.FindString(line)) + } + } + if err := scanner.Err(); err != nil { + _, _ = fmt.Fprintf(w, "slave log scan error: %v\n", err) + } +} + +type foregroundLockedWriter struct { + mu sync.Mutex + w io.Writer +} + +func (w *foregroundLockedWriter) Write(p []byte) (int, error) { + w.mu.Lock() + defer w.mu.Unlock() + return w.w.Write(p) +} diff --git a/internal/headless/slave_test.go b/internal/headless/slave_test.go new file mode 100644 index 0000000..d454856 --- /dev/null +++ b/internal/headless/slave_test.go @@ -0,0 +1,223 @@ +package headless + +import ( + "bytes" + "context" + "errors" + "fmt" + "os" + "path/filepath" + "strings" + "testing" + "time" + + "github.com/agentserver/agentserver-pkg/internal/paths" + "github.com/agentserver/agentserver-pkg/internal/terminalauth" +) + +func TestRunSlaveCreatesStableRegistryEntryAndWritesConfig(t *testing.T) { + temp := t.TempDir() + repo := filepath.Join(temp, "repo") + if err := os.Mkdir(repo, 0o755); err != nil { + t.Fatal(err) + } + authURL := "https://agent.cs.ac.cn/device?user_code=ABCD" + runner := &fakeSlaveProcessRunner{authURL: authURL} + var out bytes.Buffer + + err := RunSlave(context.Background(), SlaveOptions{ + Paths: testSlavePaths(temp), + Package: Package{SlaveAgent: filepath.Join(temp, "pkg", "slave-agent")}, + WorkDir: repo, + ComputerName: "host", + NamePrompt: func(defaultName string) (string, error) { + if defaultName != "repo" { + t.Fatalf("defaultName=%q, want repo", defaultName) + } + return "worker", nil + }, + Runner: runner, + Stdout: &out, + QR: markerQR("fake QR marker"), + CodexBin: filepath.Join(temp, "codex"), + }) + if err != nil { + t.Fatalf("RunSlave: %v", err) + } + + configBody, err := os.ReadFile(runner.request.ConfigPath) + if err != nil { + t.Fatal(err) + } + config := string(configBody) + if !strings.Contains(config, "display_name: host-worker") { + t.Fatalf("config missing display name:\n%s", config) + } + if !strings.Contains(config, fmt.Sprintf("workdir: %s", repo)) { + t.Fatalf("config missing workdir %q:\n%s", repo, config) + } + if runner.request.Exe != filepath.Join(temp, "pkg", "slave-agent") { + t.Fatalf("Exe=%q", runner.request.Exe) + } + if runner.request.WorkDir != filepath.Dir(runner.request.ConfigPath) { + t.Fatalf("WorkDir=%q, want config dir %q", runner.request.WorkDir, filepath.Dir(runner.request.ConfigPath)) + } + + gotOutput := out.String() + if !strings.Contains(gotOutput, "Agentserver 登录") { + t.Fatalf("output missing title:\n%s", gotOutput) + } + if !strings.Contains(gotOutput, authURL) { + t.Fatalf("output missing auth URL:\n%s", gotOutput) + } + if !strings.Contains(gotOutput, "fake QR marker") { + t.Fatalf("output missing QR marker:\n%s", gotOutput) + } +} + +func TestRunSlaveReusesExistingEntryWithoutPrompt(t *testing.T) { + temp := t.TempDir() + repo := filepath.Join(temp, "repo") + if err := os.Mkdir(repo, 0o755); err != nil { + t.Fatal(err) + } + p := testSlavePaths(temp) + promptCalls := 0 + firstRunner := &fakeSlaveProcessRunner{} + + if err := RunSlave(context.Background(), SlaveOptions{ + Paths: p, + Package: Package{SlaveAgent: filepath.Join(temp, "pkg", "slave-agent")}, + WorkDir: repo, + ComputerName: "host", + NamePrompt: func(string) (string, error) { + promptCalls++ + return "worker", nil + }, + Runner: firstRunner, + Stdout: ioDiscard{}, + }); err != nil { + t.Fatalf("first RunSlave: %v", err) + } + if promptCalls != 1 { + t.Fatalf("promptCalls after first run=%d, want 1", promptCalls) + } + + secondRunner := &fakeSlaveProcessRunner{} + if err := RunSlave(context.Background(), SlaveOptions{ + Paths: p, + Package: Package{SlaveAgent: filepath.Join(temp, "pkg", "slave-agent")}, + WorkDir: repo, + ComputerName: "host", + NamePrompt: func(string) (string, error) { + t.Fatal("NamePrompt called for existing folder") + return "", nil + }, + Runner: secondRunner, + Stdout: ioDiscard{}, + }); err != nil { + t.Fatalf("second RunSlave: %v", err) + } + + if firstRunner.request.ConfigPath == "" { + t.Fatal("first runner config path empty") + } + if secondRunner.request.ConfigPath != firstRunner.request.ConfigPath { + t.Fatalf("second config path=%q, want reused %q", secondRunner.request.ConfigPath, firstRunner.request.ConfigPath) + } + if promptCalls != 1 { + t.Fatalf("promptCalls after second run=%d, want 1", promptCalls) + } +} + +func TestRunSlaveStopsChildOnContextCancel(t *testing.T) { + temp := t.TempDir() + repo := filepath.Join(temp, "repo") + if err := os.Mkdir(repo, 0o755); err != nil { + t.Fatal(err) + } + ctx, cancel := context.WithCancel(context.Background()) + runner := &fakeSlaveProcessRunner{waitForCancel: true} + started := make(chan struct{}) + runner.onRun = func() { + close(started) + } + + errCh := make(chan error, 1) + go func() { + errCh <- RunSlave(ctx, SlaveOptions{ + Paths: testSlavePaths(temp), + Package: Package{SlaveAgent: filepath.Join(temp, "pkg", "slave-agent")}, + WorkDir: repo, + ComputerName: "host", + NamePrompt: func(string) (string, error) { return "worker", nil }, + Runner: runner, + Stdout: ioDiscard{}, + }) + }() + + select { + case <-started: + case <-time.After(time.Second): + t.Fatal("RunSlave did not enter runner") + } + cancel() + + var err error + select { + case err = <-errCh: + case <-time.After(time.Second): + t.Fatal("RunSlave did not stop after context cancellation") + } + if !errors.Is(err, context.Canceled) { + t.Fatalf("RunSlave error=%v, want context.Canceled", err) + } + if !runner.stopped { + t.Fatal("runner did not observe cancellation") + } +} + +type fakeSlaveProcessRunner struct { + authURL string + waitForCancel bool + request SlaveProcessRequest + stopped bool + onRun func() +} + +func (f *fakeSlaveProcessRunner) Run(ctx context.Context, req SlaveProcessRequest) error { + f.request = req + if f.onRun != nil { + f.onRun() + } + if f.authURL != "" { + req.AuthURL(f.authURL) + } + if f.waitForCancel { + <-ctx.Done() + f.stopped = ctx.Err() != nil + return ctx.Err() + } + return nil +} + +func testSlavePaths(root string) paths.Paths { + state := filepath.Join(root, "state") + return paths.Paths{ + MachineFile: filepath.Join(state, "machine.json"), + SlavesFile: filepath.Join(state, "slaves.json"), + SlavesDir: filepath.Join(state, "slaves"), + } +} + +func markerQR(marker string) terminalauth.QRWriter { + return func(w interface{ Write([]byte) (int, error) }, _ string) { + _, _ = w.Write([]byte(marker + "\n")) + } +} + +type ioDiscard struct{} + +func (ioDiscard) Write(p []byte) (int, error) { + return len(p), nil +} From a0f9c94aaca03daa2cfd99785d893c56cea07582 Mon Sep 17 00:00:00 2001 From: Zishu Yu Date: Sun, 14 Jun 2026 01:22:34 +0800 Subject: [PATCH 21/40] fix(headless): harden foreground slave runner --- internal/headless/slave.go | 136 ++++++++++++++++++---- internal/headless/slave_process_other.go | 17 +++ internal/headless/slave_process_unix.go | 29 +++++ internal/headless/slave_test.go | 139 +++++++++++++++++++++++ 4 files changed, 296 insertions(+), 25 deletions(-) create mode 100644 internal/headless/slave_process_other.go create mode 100644 internal/headless/slave_process_unix.go diff --git a/internal/headless/slave.go b/internal/headless/slave.go index 138baa6..884d0c4 100644 --- a/internal/headless/slave.go +++ b/internal/headless/slave.go @@ -3,6 +3,7 @@ package headless import ( "bufio" "context" + "errors" "fmt" "io" "os" @@ -11,6 +12,7 @@ import ( "regexp" "strings" "sync" + "time" "github.com/agentserver/agentserver-pkg/internal/paths" "github.com/agentserver/agentserver-pkg/internal/slave" @@ -47,9 +49,10 @@ func RunSlave(ctx context.Context, opts SlaveOptions) error { if stdout == nil { stdout = os.Stdout } + output := newForegroundOutput(stdout) runner := opts.Runner if runner == nil { - runner = execSlaveRunner{stdout: stdout} + runner = execSlaveRunner{output: output} } computerName := strings.TrimSpace(opts.ComputerName) if computerName == "" { @@ -96,25 +99,39 @@ func RunSlave(ctx context.Context, opts SlaveOptions) error { return err } + seenAuthURLs := map[string]struct{}{} return runner.Run(ctx, SlaveProcessRequest{ Exe: opts.Package.SlaveAgent, ConfigPath: sl.ConfigPath, WorkDir: filepath.Dir(sl.ConfigPath), AuthURL: func(url string) { - terminalauth.PrintURL(stdout, "Agentserver 登录", url, "", opts.QR) + url = sanitizeForegroundAuthURL(url) + if url == "" { + return + } + output.WithLock(func(w io.Writer) { + if _, ok := seenAuthURLs[url]; ok { + return + } + seenAuthURLs[url] = struct{}{} + terminalauth.PrintURL(w, "Agentserver 登录", url, "", opts.QR) + }) }, }) } type execSlaveRunner struct { stdout io.Writer + output *foregroundOutput } var foregroundAuthURLPattern = regexp.MustCompile(`(?i)https?://\S*(?:device|user[_-]code|verification)\S*`) +var foregroundHTTPURLPattern = regexp.MustCompile(`(?i)https?://\S+`) func (r execSlaveRunner) Run(ctx context.Context, req SlaveProcessRequest) error { - cmd := exec.CommandContext(ctx, req.Exe, req.ConfigPath) + cmd := exec.Command(req.Exe, req.ConfigPath) cmd.Dir = req.WorkDir + configureSlaveProcess(cmd) stdout, err := cmd.StdoutPipe() if err != nil { @@ -128,23 +145,35 @@ func (r execSlaveRunner) Run(ctx context.Context, req SlaveProcessRequest) error return fmt.Errorf("start slave-agent: %w", err) } - writer := r.stdout - if writer == nil { - writer = os.Stdout - } - writer = &foregroundLockedWriter{w: writer} + output := r.foregroundOutput() var wg sync.WaitGroup wg.Add(2) go func() { defer wg.Done() - copyForegroundSlaveOutput(stdout, writer, req.AuthURL) + copyForegroundSlaveOutput(stdout, output, req.AuthURL) }() go func() { defer wg.Done() - copyForegroundSlaveOutput(stderr, writer, req.AuthURL) + copyForegroundSlaveOutput(stderr, output, req.AuthURL) }() - err = cmd.Wait() + done := make(chan error, 1) + go func() { + done <- cmd.Wait() + }() + select { + case err = <-done: + case <-ctx.Done(): + terminateSlaveProcess(cmd) + select { + case <-done: + case <-time.After(2 * time.Second): + killSlaveProcess(cmd) + <-done + } + wg.Wait() + return ctx.Err() + } wg.Wait() if ctx.Err() != nil { return ctx.Err() @@ -155,28 +184,85 @@ func (r execSlaveRunner) Run(ctx context.Context, req SlaveProcessRequest) error return nil } +func (r execSlaveRunner) foregroundOutput() *foregroundOutput { + if r.output != nil { + return r.output + } + writer := r.stdout + if writer == nil { + writer = os.Stdout + } + return newForegroundOutput(writer) +} + +const maxForegroundAuthDetectionBytes = 1 << 20 + func copyForegroundSlaveOutput(r io.Reader, w io.Writer, authURL func(string)) { - scanner := bufio.NewScanner(r) - scanner.Buffer(make([]byte, 1<<16), 1<<20) - for scanner.Scan() { - line := scanner.Text() - _, _ = fmt.Fprintln(w, line) - if authURL != nil && foregroundAuthURLPattern.MatchString(line) { - authURL(foregroundAuthURLPattern.FindString(line)) + reader := bufio.NewReaderSize(r, 64*1024) + line := make([]byte, 0, 4096) + flushLine := func() { + if authURL != nil { + if match := foregroundAuthURLPattern.Find(line); len(match) > 0 { + authURL(string(match)) + } } + line = line[:0] } - if err := scanner.Err(); err != nil { - _, _ = fmt.Fprintf(w, "slave log scan error: %v\n", err) + for { + fragment, err := reader.ReadSlice('\n') + if len(fragment) > 0 { + _, _ = w.Write(fragment) + remaining := maxForegroundAuthDetectionBytes - len(line) + if remaining > 0 { + if len(fragment) > remaining { + line = append(line, fragment[:remaining]...) + } else { + line = append(line, fragment...) + } + } + } + switch { + case err == nil: + flushLine() + case errors.Is(err, bufio.ErrBufferFull): + continue + case errors.Is(err, io.EOF): + if len(line) > 0 { + flushLine() + } + return + default: + _, _ = fmt.Fprintf(w, "slave log read error: %v\n", err) + return + } } } -type foregroundLockedWriter struct { +func sanitizeForegroundAuthURL(raw string) string { + raw = strings.TrimSpace(raw) + if raw == "" { + return "" + } + return foregroundHTTPURLPattern.FindString(raw) +} + +type foregroundOutput struct { mu sync.Mutex w io.Writer } -func (w *foregroundLockedWriter) Write(p []byte) (int, error) { - w.mu.Lock() - defer w.mu.Unlock() - return w.w.Write(p) +func newForegroundOutput(w io.Writer) *foregroundOutput { + return &foregroundOutput{w: w} +} + +func (o *foregroundOutput) Write(p []byte) (int, error) { + o.mu.Lock() + defer o.mu.Unlock() + return o.w.Write(p) +} + +func (o *foregroundOutput) WithLock(fn func(io.Writer)) { + o.mu.Lock() + defer o.mu.Unlock() + fn(o.w) } diff --git a/internal/headless/slave_process_other.go b/internal/headless/slave_process_other.go new file mode 100644 index 0000000..8db6144 --- /dev/null +++ b/internal/headless/slave_process_other.go @@ -0,0 +1,17 @@ +//go:build !unix + +package headless + +import "os/exec" + +func configureSlaveProcess(_ *exec.Cmd) {} + +func terminateSlaveProcess(cmd *exec.Cmd) { + killSlaveProcess(cmd) +} + +func killSlaveProcess(cmd *exec.Cmd) { + if cmd.Process != nil { + _ = cmd.Process.Kill() + } +} diff --git a/internal/headless/slave_process_unix.go b/internal/headless/slave_process_unix.go new file mode 100644 index 0000000..1d4d9a0 --- /dev/null +++ b/internal/headless/slave_process_unix.go @@ -0,0 +1,29 @@ +//go:build unix + +package headless + +import ( + "os/exec" + "syscall" +) + +func configureSlaveProcess(cmd *exec.Cmd) { + cmd.SysProcAttr = &syscall.SysProcAttr{Setpgid: true} +} + +func terminateSlaveProcess(cmd *exec.Cmd) { + signalSlaveProcessGroup(cmd, syscall.SIGTERM) +} + +func killSlaveProcess(cmd *exec.Cmd) { + signalSlaveProcessGroup(cmd, syscall.SIGKILL) +} + +func signalSlaveProcessGroup(cmd *exec.Cmd, sig syscall.Signal) { + if cmd.Process == nil || cmd.Process.Pid <= 0 { + return + } + if err := syscall.Kill(-cmd.Process.Pid, sig); err != nil { + _ = cmd.Process.Signal(sig) + } +} diff --git a/internal/headless/slave_test.go b/internal/headless/slave_test.go index d454856..b3ac389 100644 --- a/internal/headless/slave_test.go +++ b/internal/headless/slave_test.go @@ -6,7 +6,9 @@ import ( "errors" "fmt" "os" + "os/exec" "path/filepath" + "runtime" "strings" "testing" "time" @@ -130,6 +132,39 @@ func TestRunSlaveReusesExistingEntryWithoutPrompt(t *testing.T) { } } +func TestRunSlavePrintsDuplicateAuthURLOnce(t *testing.T) { + temp := t.TempDir() + repo := filepath.Join(temp, "repo") + if err := os.Mkdir(repo, 0o755); err != nil { + t.Fatal(err) + } + authURL := "https://agent.cs.ac.cn/device?user_code=ABCD" + runner := &fakeSlaveProcessRunner{authURLs: []string{authURL, authURL}} + var out bytes.Buffer + + err := RunSlave(context.Background(), SlaveOptions{ + Paths: testSlavePaths(temp), + Package: Package{SlaveAgent: filepath.Join(temp, "pkg", "slave-agent")}, + WorkDir: repo, + ComputerName: "host", + NamePrompt: func(string) (string, error) { return "worker", nil }, + Runner: runner, + Stdout: &out, + QR: markerQR("fake QR marker"), + }) + if err != nil { + t.Fatalf("RunSlave: %v", err) + } + + gotOutput := out.String() + if got := strings.Count(gotOutput, authURL); got != 1 { + t.Fatalf("auth URL printed %d times, want 1:\n%s", got, gotOutput) + } + if got := strings.Count(gotOutput, "fake QR marker"); got != 1 { + t.Fatalf("QR marker printed %d times, want 1:\n%s", got, gotOutput) + } +} + func TestRunSlaveStopsChildOnContextCancel(t *testing.T) { temp := t.TempDir() repo := filepath.Join(temp, "repo") @@ -177,8 +212,93 @@ func TestRunSlaveStopsChildOnContextCancel(t *testing.T) { } } +func TestExecSlaveRunnerForwardsOutputAuthURLAndCancels(t *testing.T) { + if runtime.GOOS == "windows" { + t.Skip("shell script test requires /bin/sh") + } + sh, err := exec.LookPath("sh") + if err != nil { + t.Skip("sh not found") + } + temp := t.TempDir() + ready := filepath.Join(temp, "ready") + authURL := "https://agent.cs.ac.cn/device?user_code=ABCD" + script := filepath.Join(temp, "slave-agent.sh") + if err := os.WriteFile(script, []byte(fmt.Sprintf(`#!/bin/sh +echo stdout-line +echo stderr-line >&2 +echo %s +echo %s +touch %s +sleep 30 +`, authURL, authURL, ready)), 0o755); err != nil { + t.Fatal(err) + } + + ctx, cancel := context.WithCancel(context.Background()) + var out bytes.Buffer + var authURLs []string + errCh := make(chan error, 1) + go func() { + errCh <- execSlaveRunner{stdout: &out}.Run(ctx, SlaveProcessRequest{ + Exe: sh, + ConfigPath: script, + WorkDir: temp, + AuthURL: func(url string) { + authURLs = append(authURLs, url) + }, + }) + }() + + waitForFile(t, ready, time.Second) + cancel() + + select { + case err := <-errCh: + if !errors.Is(err, context.Canceled) { + t.Fatalf("Run error=%v, want context.Canceled", err) + } + case <-time.After(2 * time.Second): + t.Fatal("runner did not stop after context cancellation") + } + + gotOutput := out.String() + if !strings.Contains(gotOutput, "stdout-line") { + t.Fatalf("output missing stdout:\n%s", gotOutput) + } + if !strings.Contains(gotOutput, "stderr-line") { + t.Fatalf("output missing stderr:\n%s", gotOutput) + } + if len(authURLs) == 0 { + t.Fatalf("auth callback not observed; output:\n%s", gotOutput) + } +} + +func TestCopyForegroundSlaveOutputDrainsLongLinesAndDetectsLaterAuthURL(t *testing.T) { + authURL := "https://agent.cs.ac.cn/device?user_code=ABCD" + input := strings.Repeat("x", 1<<20+1) + "\n" + authURL + "\n" + var out bytes.Buffer + var authURLs []string + + copyForegroundSlaveOutput(strings.NewReader(input), &out, func(url string) { + authURLs = append(authURLs, url) + }) + + gotOutput := out.String() + if !strings.Contains(gotOutput, strings.Repeat("x", 128)) { + t.Fatalf("output missing long line content") + } + if !strings.Contains(gotOutput, authURL) { + t.Fatalf("output missing auth URL after long line") + } + if len(authURLs) != 1 || authURLs[0] != authURL { + t.Fatalf("authURLs=%v, want [%q]", authURLs, authURL) + } +} + type fakeSlaveProcessRunner struct { authURL string + authURLs []string waitForCancel bool request SlaveProcessRequest stopped bool @@ -193,6 +313,9 @@ func (f *fakeSlaveProcessRunner) Run(ctx context.Context, req SlaveProcessReques if f.authURL != "" { req.AuthURL(f.authURL) } + for _, authURL := range f.authURLs { + req.AuthURL(authURL) + } if f.waitForCancel { <-ctx.Done() f.stopped = ctx.Err() != nil @@ -221,3 +344,19 @@ type ioDiscard struct{} func (ioDiscard) Write(p []byte) (int, error) { return len(p), nil } + +func waitForFile(t *testing.T, path string, timeout time.Duration) { + t.Helper() + deadline := time.Now().Add(timeout) + for { + if _, err := os.Stat(path); err == nil { + return + } else if !os.IsNotExist(err) { + t.Fatal(err) + } + if time.Now().After(deadline) { + t.Fatalf("timed out waiting for %s", path) + } + time.Sleep(10 * time.Millisecond) + } +} From 6a2b62068a52a690b4329c17041e706ca44da9b9 Mon Sep 17 00:00:00 2001 From: Zishu Yu Date: Sun, 14 Jun 2026 01:30:55 +0800 Subject: [PATCH 22/40] fix(headless): validate foreground auth URLs --- internal/headless/slave.go | 39 +++++++++++++++++++--- internal/headless/slave_test.go | 59 +++++++++++++++++++++++++++++++++ 2 files changed, 94 insertions(+), 4 deletions(-) diff --git a/internal/headless/slave.go b/internal/headless/slave.go index 884d0c4..0388f20 100644 --- a/internal/headless/slave.go +++ b/internal/headless/slave.go @@ -6,6 +6,7 @@ import ( "errors" "fmt" "io" + "net/url" "os" "os/exec" "path/filepath" @@ -125,7 +126,6 @@ type execSlaveRunner struct { output *foregroundOutput } -var foregroundAuthURLPattern = regexp.MustCompile(`(?i)https?://\S*(?:device|user[_-]code|verification)\S*`) var foregroundHTTPURLPattern = regexp.MustCompile(`(?i)https?://\S+`) func (r execSlaveRunner) Run(ctx context.Context, req SlaveProcessRequest) error { @@ -202,8 +202,11 @@ func copyForegroundSlaveOutput(r io.Reader, w io.Writer, authURL func(string)) { line := make([]byte, 0, 4096) flushLine := func() { if authURL != nil { - if match := foregroundAuthURLPattern.Find(line); len(match) > 0 { - authURL(string(match)) + for _, match := range foregroundHTTPURLPattern.FindAll(line, -1) { + if url := sanitizeForegroundAuthURL(string(match)); url != "" { + authURL(url) + break + } } } line = line[:0] @@ -243,7 +246,35 @@ func sanitizeForegroundAuthURL(raw string) string { if raw == "" { return "" } - return foregroundHTTPURLPattern.FindString(raw) + for _, match := range foregroundHTTPURLPattern.FindAllString(raw, -1) { + if isForegroundAuthURL(match) { + return match + } + } + return "" +} + +func isForegroundAuthURL(raw string) bool { + parsed, err := url.Parse(raw) + if err != nil { + return false + } + switch strings.ToLower(parsed.Scheme) { + case "http", "https": + default: + return false + } + if !strings.EqualFold(parsed.Hostname(), "agent.cs.ac.cn") { + return false + } + path := strings.ToLower(parsed.EscapedPath()) + query := strings.ToLower(parsed.RawQuery) + for _, marker := range []string{"device", "user_code", "user-code", "verification"} { + if strings.Contains(path, marker) || strings.Contains(query, marker) { + return true + } + } + return false } type foregroundOutput struct { diff --git a/internal/headless/slave_test.go b/internal/headless/slave_test.go index b3ac389..72789c6 100644 --- a/internal/headless/slave_test.go +++ b/internal/headless/slave_test.go @@ -165,6 +165,43 @@ func TestRunSlavePrintsDuplicateAuthURLOnce(t *testing.T) { } } +func TestRunSlaveIgnoresNonAgentserverAuthURLs(t *testing.T) { + temp := t.TempDir() + repo := filepath.Join(temp, "repo") + if err := os.Mkdir(repo, 0o755); err != nil { + t.Fatal(err) + } + runner := &fakeSlaveProcessRunner{authURLs: []string{ + "https://example.com/help", + "https://example.com/device-docs", + "ftp://agent.cs.ac.cn/device?user_code=ABCD", + "https://agent.cs.ac.cn/help", + }} + var out bytes.Buffer + + err := RunSlave(context.Background(), SlaveOptions{ + Paths: testSlavePaths(temp), + Package: Package{SlaveAgent: filepath.Join(temp, "pkg", "slave-agent")}, + WorkDir: repo, + ComputerName: "host", + NamePrompt: func(string) (string, error) { return "worker", nil }, + Runner: runner, + Stdout: &out, + QR: markerQR("fake QR marker"), + }) + if err != nil { + t.Fatalf("RunSlave: %v", err) + } + + gotOutput := out.String() + if strings.Contains(gotOutput, "Agentserver 登录") { + t.Fatalf("output printed auth title for unrelated URLs:\n%s", gotOutput) + } + if strings.Contains(gotOutput, "fake QR marker") { + t.Fatalf("output printed QR for unrelated URLs:\n%s", gotOutput) + } +} + func TestRunSlaveStopsChildOnContextCancel(t *testing.T) { temp := t.TempDir() repo := filepath.Join(temp, "repo") @@ -296,6 +333,28 @@ func TestCopyForegroundSlaveOutputDrainsLongLinesAndDetectsLaterAuthURL(t *testi } } +func TestCopyForegroundSlaveOutputIgnoresNonAgentserverAuthURLs(t *testing.T) { + authURL := "https://agent.cs.ac.cn/device?user_code=ABCD" + input := strings.Join([]string{ + "https://example.com/help", + "https://example.com/device-docs", + authURL, + }, "\n") + var out bytes.Buffer + var authURLs []string + + copyForegroundSlaveOutput(strings.NewReader(input), &out, func(url string) { + authURLs = append(authURLs, url) + }) + + if gotOutput := out.String(); !strings.Contains(gotOutput, authURL) { + t.Fatalf("output missing accepted auth URL:\n%s", gotOutput) + } + if len(authURLs) != 1 || authURLs[0] != authURL { + t.Fatalf("authURLs=%v, want [%q]", authURLs, authURL) + } +} + type fakeSlaveProcessRunner struct { authURL string authURLs []string From 3ebfc5b1b1bb523d65f190b13c7e86d5bf9bd07b Mon Sep 17 00:00:00 2001 From: Zishu Yu Date: Sun, 14 Jun 2026 01:40:48 +0800 Subject: [PATCH 23/40] fix(headless): tighten foreground auth detection --- internal/headless/slave.go | 82 ++++++++++++++++++++++++---- internal/headless/slave_test.go | 95 +++++++++++++++++++++++++++++++++ 2 files changed, 168 insertions(+), 9 deletions(-) diff --git a/internal/headless/slave.go b/internal/headless/slave.go index 0388f20..d5875bb 100644 --- a/internal/headless/slave.go +++ b/internal/headless/slave.go @@ -146,15 +146,16 @@ func (r execSlaveRunner) Run(ctx context.Context, req SlaveProcessRequest) error } output := r.foregroundOutput() + authURL := serializeAuthURLCallback(req.AuthURL) var wg sync.WaitGroup wg.Add(2) go func() { defer wg.Done() - copyForegroundSlaveOutput(stdout, output, req.AuthURL) + copyForegroundSlaveOutput(stdout, output, authURL) }() go func() { defer wg.Done() - copyForegroundSlaveOutput(stderr, output, req.AuthURL) + copyForegroundSlaveOutput(stderr, output, authURL) }() done := make(chan error, 1) @@ -241,19 +242,60 @@ func copyForegroundSlaveOutput(r io.Reader, w io.Writer, authURL func(string)) { } } +func serializeAuthURLCallback(authURL func(string)) func(string) { + if authURL == nil { + return nil + } + var mu sync.Mutex + return func(url string) { + mu.Lock() + defer mu.Unlock() + authURL(url) + } +} + func sanitizeForegroundAuthURL(raw string) string { raw = strings.TrimSpace(raw) if raw == "" { return "" } for _, match := range foregroundHTTPURLPattern.FindAllString(raw, -1) { - if isForegroundAuthURL(match) { - return match + candidate := trimForegroundURLCandidate(match) + if isForegroundAuthURL(candidate) { + return candidate } } return "" } +const foregroundURLTrailingDelimiters = ".,;:!?)]}>\"'`" + +func trimForegroundURLCandidate(raw string) string { + candidate := strings.TrimSpace(raw) + for { + before := candidate + candidate = trimTrailingANSIReset(candidate) + candidate = strings.TrimRight(candidate, foregroundURLTrailingDelimiters) + candidate = trimTrailingANSIReset(candidate) + if candidate == before { + return candidate + } + } +} + +func trimTrailingANSIReset(raw string) string { + for { + switch { + case strings.HasSuffix(raw, "\x1b[0m"): + raw = strings.TrimSuffix(raw, "\x1b[0m") + case strings.HasSuffix(raw, "\x1b[m"): + raw = strings.TrimSuffix(raw, "\x1b[m") + default: + return raw + } + } +} + func isForegroundAuthURL(raw string) bool { parsed, err := url.Parse(raw) if err != nil { @@ -264,19 +306,41 @@ func isForegroundAuthURL(raw string) bool { default: return false } - if !strings.EqualFold(parsed.Hostname(), "agent.cs.ac.cn") { + if !strings.EqualFold(parsed.Hostname(), "agent.cs.ac.cn") || parsed.Port() != "" { return false } - path := strings.ToLower(parsed.EscapedPath()) - query := strings.ToLower(parsed.RawQuery) - for _, marker := range []string{"device", "user_code", "user-code", "verification"} { - if strings.Contains(path, marker) || strings.Contains(query, marker) { + hasCodeQuery := hasForegroundAuthCodeQuery(parsed.Query()) + hasUserCodePath := hasForegroundAuthPathSegment(parsed.Path, "user-code", "user_code") + if hasUserCodePath { + return true + } + if hasCodeQuery && hasForegroundAuthPathSegment(parsed.Path, "device", "verification") { + return true + } + return false +} + +func hasForegroundAuthCodeQuery(query url.Values) bool { + for key := range query { + if strings.EqualFold(key, "user_code") || strings.EqualFold(key, "code") { return true } } return false } +func hasForegroundAuthPathSegment(path string, markers ...string) bool { + for _, segment := range strings.Split(path, "/") { + segment = strings.ToLower(segment) + for _, marker := range markers { + if strings.Contains(segment, marker) { + return true + } + } + } + return false +} + type foregroundOutput struct { mu sync.Mutex w io.Writer diff --git a/internal/headless/slave_test.go b/internal/headless/slave_test.go index 72789c6..f7230c5 100644 --- a/internal/headless/slave_test.go +++ b/internal/headless/slave_test.go @@ -10,6 +10,7 @@ import ( "path/filepath" "runtime" "strings" + "sync" "testing" "time" @@ -311,6 +312,61 @@ sleep 30 } } +func TestExecSlaveRunnerSerializesAuthCallbacks(t *testing.T) { + if runtime.GOOS == "windows" { + t.Skip("shell script test requires /bin/sh") + } + sh, err := exec.LookPath("sh") + if err != nil { + t.Skip("sh not found") + } + temp := t.TempDir() + authURL := "https://agent.cs.ac.cn/device?user_code=ABCD" + script := filepath.Join(temp, "slave-agent.sh") + if err := os.WriteFile(script, []byte(fmt.Sprintf(`#!/bin/sh +for i in 1 2 3 4 5; do + echo %s + echo %s >&2 +done +`, authURL, authURL)), 0o755); err != nil { + t.Fatal(err) + } + + var ( + mu sync.Mutex + inAuth bool + authURLs []string + ) + err = execSlaveRunner{stdout: ioDiscard{}}.Run(context.Background(), SlaveProcessRequest{ + Exe: sh, + ConfigPath: script, + WorkDir: temp, + AuthURL: func(url string) { + mu.Lock() + if inAuth { + mu.Unlock() + t.Errorf("AuthURL called concurrently") + return + } + inAuth = true + mu.Unlock() + + time.Sleep(2 * time.Millisecond) + authURLs = append(authURLs, url) + + mu.Lock() + inAuth = false + mu.Unlock() + }, + }) + if err != nil { + t.Fatalf("Run: %v", err) + } + if len(authURLs) == 0 { + t.Fatal("auth callback not observed") + } +} + func TestCopyForegroundSlaveOutputDrainsLongLinesAndDetectsLaterAuthURL(t *testing.T) { authURL := "https://agent.cs.ac.cn/device?user_code=ABCD" input := strings.Repeat("x", 1<<20+1) + "\n" + authURL + "\n" @@ -333,11 +389,50 @@ func TestCopyForegroundSlaveOutputDrainsLongLinesAndDetectsLaterAuthURL(t *testi } } +func TestSanitizeForegroundAuthURL(t *testing.T) { + tests := []struct { + name string + raw string + want string + }{ + { + name: "trims trailing period", + raw: "open https://agent.cs.ac.cn/device?user_code=ABCD.", + want: "https://agent.cs.ac.cn/device?user_code=ABCD", + }, + { + name: "trims trailing closing paren", + raw: "https://agent.cs.ac.cn/verification?code=ABCD)", + want: "https://agent.cs.ac.cn/verification?code=ABCD", + }, + { + name: "rejects marker in unrelated query", + raw: "https://agent.cs.ac.cn/help?topic=device", + want: "", + }, + { + name: "rejects marker in unrelated path segment", + raw: "https://agent.cs.ac.cn/device-docs", + want: "", + }, + } + + for _, tt := range tests { + t.Run(tt.name, func(t *testing.T) { + if got := sanitizeForegroundAuthURL(tt.raw); got != tt.want { + t.Fatalf("sanitizeForegroundAuthURL(%q)=%q, want %q", tt.raw, got, tt.want) + } + }) + } +} + func TestCopyForegroundSlaveOutputIgnoresNonAgentserverAuthURLs(t *testing.T) { authURL := "https://agent.cs.ac.cn/device?user_code=ABCD" input := strings.Join([]string{ "https://example.com/help", "https://example.com/device-docs", + "https://agent.cs.ac.cn/help?topic=device", + "https://agent.cs.ac.cn/device-docs", authURL, }, "\n") var out bytes.Buffer From fcba0ebf2bcb334b7d9975d0579e106422f014ff Mon Sep 17 00:00:00 2001 From: Zishu Yu Date: Sun, 14 Jun 2026 01:43:00 +0800 Subject: [PATCH 24/40] fix(headless): require exact auth path markers --- internal/headless/slave.go | 2 +- internal/headless/slave_test.go | 5 +++++ 2 files changed, 6 insertions(+), 1 deletion(-) diff --git a/internal/headless/slave.go b/internal/headless/slave.go index d5875bb..4403d82 100644 --- a/internal/headless/slave.go +++ b/internal/headless/slave.go @@ -333,7 +333,7 @@ func hasForegroundAuthPathSegment(path string, markers ...string) bool { for _, segment := range strings.Split(path, "/") { segment = strings.ToLower(segment) for _, marker := range markers { - if strings.Contains(segment, marker) { + if segment == marker { return true } } diff --git a/internal/headless/slave_test.go b/internal/headless/slave_test.go index f7230c5..2bcf8ae 100644 --- a/internal/headless/slave_test.go +++ b/internal/headless/slave_test.go @@ -415,6 +415,11 @@ func TestSanitizeForegroundAuthURL(t *testing.T) { raw: "https://agent.cs.ac.cn/device-docs", want: "", }, + { + name: "rejects unrelated path segment with auth query", + raw: "https://agent.cs.ac.cn/device-docs?user_code=ABCD", + want: "", + }, } for _, tt := range tests { From 6fa97cb71c2ff63ff4af7f1c261728b4a7ff1792 Mon Sep 17 00:00:00 2001 From: Zishu Yu Date: Sun, 14 Jun 2026 02:02:34 +0800 Subject: [PATCH 25/40] feat(headless): install Linux driver MCP --- internal/headless/driver.go | 369 ++++++++++++++++++++++++++ internal/headless/driver_test.go | 428 +++++++++++++++++++++++++++++++ 2 files changed, 797 insertions(+) create mode 100644 internal/headless/driver.go create mode 100644 internal/headless/driver_test.go diff --git a/internal/headless/driver.go b/internal/headless/driver.go new file mode 100644 index 0000000..8dfde8d --- /dev/null +++ b/internal/headless/driver.go @@ -0,0 +1,369 @@ +package headless + +import ( + "context" + "errors" + "fmt" + "io" + "os" + osexec "os/exec" + "path/filepath" + "strings" + + "github.com/agentserver/agentserver-pkg/internal/agentserver" + "github.com/agentserver/agentserver-pkg/internal/codex" + "github.com/agentserver/agentserver-pkg/internal/loom" + "github.com/agentserver/agentserver-pkg/internal/oauth" + "github.com/agentserver/agentserver-pkg/internal/paths" + "github.com/agentserver/agentserver-pkg/internal/secrets" + "github.com/agentserver/agentserver-pkg/internal/state" + "github.com/agentserver/agentserver-pkg/internal/terminalauth" +) + +const ( + defaultAgentserverEndpoint = "https://agent.cs.ac.cn" + driverProxySecretKey = "agentserver_ws_api_key" + driverTunnelSecretKey = "agentserver_tunnel_token" +) + +type AgentRegistrar interface { + RegisterAgent(context.Context, string, string, string) (agentserver.AgentRegistration, error) + Whoami(context.Context, string) (agentserver.Identity, error) +} + +type DriverOptions struct { + Paths paths.Paths + Package Package + Secrets secrets.Store + ComputerName string + AS AgentRegistrar + ASOAuth oauth.Config + RequestDeviceCode func(context.Context, oauth.Config) (oauth.DeviceCodeChallenge, error) + PollToken func(context.Context, oauth.Config, oauth.DeviceCodeChallenge) (oauth.Token, error) + Stdout io.Writer + QR terminalauth.QRWriter +} + +type DriverMCPOptions struct { + Paths paths.Paths + Package Package + Secrets secrets.Store + WorkDir string + Exec func(context.Context, string, []string) error +} + +func InstallDriver(ctx context.Context, opts DriverOptions) error { + return ensureDriver(ctx, opts, false) +} + +func SwitchWorkspace(ctx context.Context, opts DriverOptions) error { + return ensureDriver(ctx, opts, true) +} + +func ServeDriverMCP(ctx context.Context, opts DriverMCPOptions) error { + if opts.Secrets == nil { + return errors.New("driver MCP: secrets store required") + } + workDir := strings.TrimSpace(opts.WorkDir) + if workDir == "" { + wd, err := os.Getwd() + if err != nil { + return fmt.Errorf("get current workdir: %w", err) + } + workDir = wd + } + st, err := state.NewStore(opts.Paths.StateFile).Load() + if err != nil { + return err + } + proxyToken, tunnelToken, err := driverSecrets(opts.Secrets) + if err != nil { + return err + } + sessionConfig := filepath.Join(opts.Paths.InstallRoot, "driver-mcp", fmt.Sprintf("driver-%d.yaml", os.Getpid())) + if err := writeDriverConfig(sessionConfig, st, proxyToken, tunnelToken, workDir); err != nil { + return err + } + run := opts.Exec + if run == nil { + run = execDriverMCP + } + return run(ctx, opts.Package.DriverAgent, []string{"serve-mcp", "--config", sessionConfig}) +} + +func ensureDriver(ctx context.Context, opts DriverOptions, forceLogin bool) error { + if opts.Secrets == nil { + return errors.New("install driver: secrets store required") + } + opts = defaultDriverOptions(opts) + + registered, err := driverRegistered(opts.Paths, opts.Secrets) + if err != nil { + return err + } + if forceLogin || !registered { + if err := registerDriver(ctx, opts); err != nil { + return err + } + } + return refreshDriverFiles(opts.Paths, opts.Package, opts.Secrets) +} + +func defaultDriverOptions(opts DriverOptions) DriverOptions { + opts.ASOAuth = normalizeASOAuth(opts.ASOAuth) + if opts.Stdout == nil { + opts.Stdout = os.Stdout + } + if opts.RequestDeviceCode == nil { + opts.RequestDeviceCode = oauth.RequestDeviceCode + } + if opts.PollToken == nil { + opts.PollToken = oauth.PollToken + } + if opts.AS == nil { + opts.AS = agentserver.New(defaultASBaseURL(opts.ASOAuth)) + } + return opts +} + +func normalizeASOAuth(cfg oauth.Config) oauth.Config { + endpoint := strings.TrimSpace(cfg.Endpoint) + if endpoint == "" { + endpoint = defaultAgentserverEndpoint + } + base := agentserver.OAuthConfig(endpoint) + if cfg.AuthPath != "" { + base.AuthPath = cfg.AuthPath + } + if cfg.TokenPath != "" { + base.TokenPath = cfg.TokenPath + } + if cfg.ClientID != "" { + base.ClientID = cfg.ClientID + } + if cfg.Scope != "" { + base.Scope = cfg.Scope + } + return base +} + +func defaultASBaseURL(cfg oauth.Config) string { + if endpoint := strings.TrimRight(strings.TrimSpace(cfg.Endpoint), "/"); endpoint != "" { + return endpoint + } + return defaultAgentserverEndpoint +} + +func registerDriver(ctx context.Context, opts DriverOptions) error { + computerName := strings.TrimSpace(opts.ComputerName) + if computerName == "" { + hostname, err := os.Hostname() + if err != nil { + return fmt.Errorf("hostname: %w", err) + } + computerName = hostname + } + ch, err := opts.RequestDeviceCode(ctx, opts.ASOAuth) + if err != nil { + return fmt.Errorf("request agentserver device code: %w", err) + } + terminalauth.PrintChallenge(opts.Stdout, "Agentserver 登录", ch, opts.QR) + tok, err := opts.PollToken(ctx, opts.ASOAuth, ch) + if err != nil { + return fmt.Errorf("poll agentserver token: %w", err) + } + reg, err := opts.AS.RegisterAgent(ctx, tok.AccessToken, computerName+"-星池指挥官", "custom") + if err != nil { + return fmt.Errorf("register agentserver agent: %w", err) + } + if reg.ProxyToken == "" || reg.TunnelToken == "" || reg.SandboxID == "" || reg.ShortID == "" { + return errors.New("register agentserver agent: incomplete registration") + } + + workspace := resolveDriverWorkspace(ctx, opts.AS, reg, tok.AccessToken) + if workspace.ID == "" { + return errors.New("resolve agentserver workspace: empty workspace id") + } + if err := opts.Secrets.Set(driverProxySecretKey, reg.ProxyToken); err != nil { + return err + } + if err := opts.Secrets.Set(driverTunnelSecretKey, reg.TunnelToken); err != nil { + return err + } + return state.NewStore(opts.Paths.StateFile).Update(func(s *state.State) error { + s.Agentserver.BaseURL = defaultASBaseURL(opts.ASOAuth) + s.Agentserver.SandboxID = reg.SandboxID + s.Agentserver.ShortID = reg.ShortID + s.Agentserver.WorkspaceID = workspace.ID + s.Agentserver.WorkspaceName = workspace.Name + s.Agentserver.WorkspaceAPIKeySuffix = lastN(reg.ProxyToken, 4) + return nil + }) +} + +func resolveDriverWorkspace(ctx context.Context, as AgentRegistrar, reg agentserver.AgentRegistration, accessToken string) agentserver.Workspace { + if as != nil { + if identity, err := as.Whoami(ctx, reg.ProxyToken); err == nil && identity.Workspace.ID != "" { + return identity.Workspace + } + } + if ws, ok := agentserver.WorkspaceFromToken(accessToken); ok { + return ws + } + return agentserver.Workspace{ID: reg.WorkspaceID} +} + +func driverRegistered(p paths.Paths, sec secrets.Store) (bool, error) { + st, err := state.NewStore(p.StateFile).Load() + if err != nil { + return false, err + } + if strings.TrimSpace(st.Agentserver.SandboxID) == "" || strings.TrimSpace(st.Agentserver.WorkspaceID) == "" { + return false, nil + } + if ok, err := secretPresent(sec, driverProxySecretKey); err != nil || !ok { + return false, err + } + if ok, err := secretPresent(sec, driverTunnelSecretKey); err != nil || !ok { + return false, err + } + return true, nil +} + +func secretPresent(sec secrets.Store, key string) (bool, error) { + value, err := sec.Get(key) + if errors.Is(err, secrets.ErrNotFound) { + return false, nil + } + if err != nil { + return false, err + } + return strings.TrimSpace(value) != "", nil +} + +func refreshDriverFiles(p paths.Paths, pkg Package, sec secrets.Store) error { + st, err := state.NewStore(p.StateFile).Load() + if err != nil { + return err + } + proxyToken, tunnelToken, err := driverSecrets(sec) + if err != nil { + return err + } + if err := writeDriverConfig(persistentDriverConfigPath(p), st, proxyToken, tunnelToken, driverUserHome(p)); err != nil { + return err + } + if err := loom.InstallDriverSupport(loom.DriverSupportInput{ + UserHome: driverUserHome(p), + SkillsArchivePath: filepath.Join(driverPackageDir(pkg), "driver-skills.tar.gz"), + SuperpowerSkillsArchivePath: filepath.Join(driverPackageDir(pkg), "driver-superpower-skills.tar.gz"), + CodexPromptsArchivePath: filepath.Join(driverPackageDir(pkg), "driver-codex-prompts.tar.gz"), + }); err != nil { + return fmt.Errorf("install loom driver support: %w", err) + } + enabled := true + if err := codex.UpdateMCPServer(p.CodexConfigFile, "driver", codex.MCPServer{ + Command: pkg.AgentserverExe, + Args: []string{"serve-driver-mcp"}, + StartupTimeoutSec: 30, + ToolTimeoutSec: 120, + Enabled: &enabled, + }); err != nil { + return fmt.Errorf("configure codex mcp driver: %w", err) + } + return nil +} + +func driverSecrets(sec secrets.Store) (string, string, error) { + proxyToken, err := sec.Get(driverProxySecretKey) + if err != nil { + return "", "", fmt.Errorf("agentserver registration missing proxy token: %w", err) + } + tunnelToken, err := sec.Get(driverTunnelSecretKey) + if err != nil { + return "", "", fmt.Errorf("agentserver registration missing tunnel token: %w", err) + } + return proxyToken, tunnelToken, nil +} + +func writeDriverConfig(path string, st *state.State, proxyToken, tunnelToken, workDir string) error { + if st == nil { + return errors.New("driver state required") + } + if err := loom.WriteDriverConfig(path, loom.DriverConfig{ + ServerURL: defaultString(st.Agentserver.BaseURL, defaultAgentserverEndpoint), + ServerName: driverServerName(st.Agentserver), + SandboxID: st.Agentserver.SandboxID, + TunnelToken: tunnelToken, + ProxyToken: proxyToken, + WorkspaceID: st.Agentserver.WorkspaceID, + WorkspaceName: st.Agentserver.WorkspaceName, + ShortID: defaultString(st.Agentserver.ShortID, st.Agentserver.SandboxID), + DisplayName: "星池指挥官", + Description: "星池指挥官本地协作驱动。", + CodexBin: "codex", + CodexWorkDir: workDir, + }); err != nil { + return fmt.Errorf("configure loom driver: %w", err) + } + return nil +} + +func execDriverMCP(ctx context.Context, exe string, args []string) error { + cmd := osexec.CommandContext(ctx, exe, args...) + cmd.Stdin = os.Stdin + cmd.Stdout = os.Stdout + cmd.Stderr = os.Stderr + if err := cmd.Run(); err != nil { + return fmt.Errorf("driver-agent exited: %w", err) + } + return nil +} + +func persistentDriverConfigPath(p paths.Paths) string { + return filepath.Join(driverUserHome(p), ".config", "multi-agent", "driver.yaml") +} + +func driverUserHome(p paths.Paths) string { + if strings.TrimSpace(p.UserHome) != "" { + return p.UserHome + } + if home, err := os.UserHomeDir(); err == nil { + return home + } + return "" +} + +func driverPackageDir(pkg Package) string { + if strings.TrimSpace(pkg.PackageDir) != "" { + return pkg.PackageDir + } + if strings.TrimSpace(pkg.DriverAgent) != "" { + return filepath.Dir(pkg.DriverAgent) + } + return filepath.Dir(pkg.AgentserverExe) +} + +func driverServerName(st state.AgentserverState) string { + if st.ShortID != "" { + return "driver-" + st.ShortID + } + if st.SandboxID != "" { + return "driver-" + st.SandboxID + } + return "driver-local" +} + +func defaultString(v, fallback string) string { + if v != "" { + return v + } + return fallback +} + +func lastN(s string, n int) string { + if n <= 0 || len(s) <= n { + return s + } + return s[len(s)-n:] +} diff --git a/internal/headless/driver_test.go b/internal/headless/driver_test.go new file mode 100644 index 0000000..f58789b --- /dev/null +++ b/internal/headless/driver_test.go @@ -0,0 +1,428 @@ +package headless + +import ( + "bytes" + "context" + "encoding/base64" + "encoding/json" + "errors" + "os" + "path/filepath" + "strings" + "testing" + + "github.com/agentserver/agentserver-pkg/internal/agentserver" + "github.com/agentserver/agentserver-pkg/internal/oauth" + "github.com/agentserver/agentserver-pkg/internal/paths" + "github.com/agentserver/agentserver-pkg/internal/secrets" + "github.com/agentserver/agentserver-pkg/internal/state" +) + +func TestInstallDriverRegistersAgentAndMCPEntrypoint(t *testing.T) { + temp := t.TempDir() + p := testDriverPaths(temp) + sec := secrets.New(p.SecretsFile) + fakeAS := &fakeDriverAgentserver{ + reg: agentserver.AgentRegistration{ + SandboxID: "sb-1", + TunnelToken: "tunnel-token", + ProxyToken: "sandbox-proxy-token", + WorkspaceID: "ws-reg", + ShortID: "abc123", + }, + whoami: agentserver.Identity{ + Workspace: agentserver.Workspace{ID: "ws-1", Name: "Workspace One"}, + }, + } + var out bytes.Buffer + + err := InstallDriver(context.Background(), DriverOptions{ + Paths: p, + Package: testDriverPackage(temp), + Secrets: sec, + ComputerName: "host", + AS: fakeAS, + ASOAuth: oauth.Config{Endpoint: "https://agent.test"}, + RequestDeviceCode: func(context.Context, oauth.Config) (oauth.DeviceCodeChallenge, error) { + return oauth.DeviceCodeChallenge{ + DeviceCode: "device", + UserCode: "AS-CODE", + VerificationURIComplete: "https://agent.test/device?user_code=AS-CODE", + }, nil + }, + PollToken: func(context.Context, oauth.Config, oauth.DeviceCodeChallenge) (oauth.Token, error) { + return oauth.Token{AccessToken: "oauth-token"}, nil + }, + Stdout: &out, + QR: driverMarkerQR("fake QR"), + }) + if err != nil { + t.Fatalf("InstallDriver: %v", err) + } + + if fakeAS.registerName != "host-星池指挥官" { + t.Fatalf("register name=%q, want host-星池指挥官", fakeAS.registerName) + } + if fakeAS.registerType != "custom" { + t.Fatalf("register type=%q, want custom", fakeAS.registerType) + } + if !strings.Contains(out.String(), "AS-CODE") || !strings.Contains(out.String(), "fake QR") { + t.Fatalf("output missing device code or QR marker:\n%s", out.String()) + } + if got, _ := sec.Get("agentserver_ws_api_key"); got != "sandbox-proxy-token" { + t.Fatalf("agentserver_ws_api_key=%q", got) + } + if got, _ := sec.Get("agentserver_tunnel_token"); got != "tunnel-token" { + t.Fatalf("agentserver_tunnel_token=%q", got) + } + + st, err := state.NewStore(p.StateFile).Load() + if err != nil { + t.Fatal(err) + } + if st.Agentserver.WorkspaceID != "ws-1" || st.Agentserver.WorkspaceName != "Workspace One" { + t.Fatalf("workspace state=%+v", st.Agentserver) + } + + config := readTextFile(t, p.CodexConfigFile) + for _, want := range []string{ + `[mcp_servers.driver]`, + `command = "` + filepath.ToSlash(testDriverPackage(temp).AgentserverExe) + `"`, + `args = ["serve-driver-mcp"]`, + `startup_timeout_sec = 30`, + `tool_timeout_sec = 120`, + `enabled = true`, + } { + if !strings.Contains(config, want) { + t.Fatalf("config missing %q:\n%s", want, config) + } + } +} + +func TestServeDriverMCPWritesSessionConfigWithCurrentWorkdir(t *testing.T) { + temp := t.TempDir() + p := testDriverPaths(temp) + pkg := testDriverPackage(temp) + sec := secrets.New(p.SecretsFile) + if err := sec.Set("agentserver_ws_api_key", "proxy-token"); err != nil { + t.Fatal(err) + } + if err := sec.Set("agentserver_tunnel_token", "tunnel-token"); err != nil { + t.Fatal(err) + } + writeDriverState(t, p, state.AgentserverState{ + BaseURL: "https://agent.test", + SandboxID: "sb-1", + ShortID: "abc123", + WorkspaceID: "ws-1", + WorkspaceName: "Workspace One", + }) + sessionDir := filepath.Join(temp, "repo") + if err := os.Mkdir(sessionDir, 0o755); err != nil { + t.Fatal(err) + } + execCalled := false + + err := ServeDriverMCP(context.Background(), DriverMCPOptions{ + Paths: p, + Package: pkg, + Secrets: sec, + WorkDir: sessionDir, + Exec: func(_ context.Context, exe string, args []string) error { + execCalled = true + if exe != pkg.DriverAgent { + t.Fatalf("exe=%q, want %q", exe, pkg.DriverAgent) + } + if len(args) != 3 || args[0] != "serve-mcp" || args[1] != "--config" { + t.Fatalf("args=%v, want serve-mcp --config ", args) + } + wantPrefix := filepath.Join(p.InstallRoot, "driver-mcp", "driver-") + if !strings.HasPrefix(args[2], wantPrefix) || filepath.Ext(args[2]) != ".yaml" { + t.Fatalf("config path=%q, want %s*.yaml", args[2], wantPrefix) + } + body := readTextFile(t, args[2]) + wantWorkdir := `workdir: "` + filepath.ToSlash(sessionDir) + `"` + if !strings.Contains(body, wantWorkdir) { + t.Fatalf("config missing %q:\n%s", wantWorkdir, body) + } + return nil + }, + }) + if err != nil { + t.Fatalf("ServeDriverMCP: %v", err) + } + if !execCalled { + t.Fatal("Exec was not called") + } +} + +func TestSwitchWorkspaceForcesDeviceLogin(t *testing.T) { + t.Run("empty state", func(t *testing.T) { + temp := t.TempDir() + p := testDriverPaths(temp) + calls := 0 + + err := SwitchWorkspace(context.Background(), DriverOptions{ + Paths: p, + Package: testDriverPackage(temp), + Secrets: secrets.New(p.SecretsFile), + ComputerName: "host", + AS: fakeRegisteredDriverAS(), + ASOAuth: oauth.Config{Endpoint: "https://agent.test"}, + RequestDeviceCode: func(context.Context, oauth.Config) (oauth.DeviceCodeChallenge, error) { + calls++ + return driverChallenge(), nil + }, + PollToken: driverPollToken, + Stdout: ioDiscard{}, + QR: driverMarkerQR("fake QR"), + }) + if err != nil { + t.Fatalf("SwitchWorkspace: %v", err) + } + if calls != 1 { + t.Fatalf("RequestDeviceCode calls=%d, want 1", calls) + } + }) + + t.Run("existing registration", func(t *testing.T) { + temp := t.TempDir() + p := testDriverPaths(temp) + sec := secrets.New(p.SecretsFile) + if err := sec.Set("agentserver_ws_api_key", "proxy-old"); err != nil { + t.Fatal(err) + } + if err := sec.Set("agentserver_tunnel_token", "tunnel-old"); err != nil { + t.Fatal(err) + } + writeDriverState(t, p, state.AgentserverState{ + BaseURL: "https://agent.old", + SandboxID: "sb-old", + ShortID: "old123", + WorkspaceID: "ws-old", + }) + calls := 0 + + err := SwitchWorkspace(context.Background(), DriverOptions{ + Paths: p, + Package: testDriverPackage(temp), + Secrets: sec, + ComputerName: "host", + AS: fakeRegisteredDriverAS(), + ASOAuth: oauth.Config{Endpoint: "https://agent.test"}, + RequestDeviceCode: func(context.Context, oauth.Config) (oauth.DeviceCodeChallenge, error) { + calls++ + return driverChallenge(), nil + }, + PollToken: driverPollToken, + Stdout: ioDiscard{}, + QR: driverMarkerQR("fake QR"), + }) + if err != nil { + t.Fatalf("SwitchWorkspace: %v", err) + } + if calls != 1 { + t.Fatalf("RequestDeviceCode calls=%d, want forced login", calls) + } + }) +} + +func TestInstallDriverSkipsDeviceFlowWhenAlreadyRegistered(t *testing.T) { + temp := t.TempDir() + p := testDriverPaths(temp) + sec := secrets.New(p.SecretsFile) + if err := sec.Set("agentserver_ws_api_key", "proxy-token"); err != nil { + t.Fatal(err) + } + if err := sec.Set("agentserver_tunnel_token", "tunnel-token"); err != nil { + t.Fatal(err) + } + writeDriverState(t, p, state.AgentserverState{ + BaseURL: "https://agent.test", + SandboxID: "sb-1", + ShortID: "abc123", + WorkspaceID: "ws-1", + }) + + err := InstallDriver(context.Background(), DriverOptions{ + Paths: p, + Package: testDriverPackage(temp), + Secrets: sec, + AS: fakeRegisteredDriverAS(), + RequestDeviceCode: func(context.Context, oauth.Config) (oauth.DeviceCodeChallenge, error) { + t.Fatal("RequestDeviceCode called for existing registration") + return oauth.DeviceCodeChallenge{}, nil + }, + PollToken: driverPollToken, + Stdout: ioDiscard{}, + }) + if err != nil { + t.Fatalf("InstallDriver: %v", err) + } + config := readTextFile(t, p.CodexConfigFile) + if !strings.Contains(config, `[mcp_servers.driver]`) { + t.Fatalf("Codex MCP config was not written:\n%s", config) + } +} + +func TestInstallDriverErrorsIfSecretsNil(t *testing.T) { + err := InstallDriver(context.Background(), DriverOptions{ + Paths: testDriverPaths(t.TempDir()), + Package: testDriverPackage(t.TempDir()), + }) + if err == nil { + t.Fatal("InstallDriver succeeded with nil Secrets") + } + if !strings.Contains(err.Error(), "secrets store required") { + t.Fatalf("error=%v, want secrets store required", err) + } +} + +func TestInstallDriverStoresWorkspaceFallbackFromAccessTokenClaim(t *testing.T) { + temp := t.TempDir() + p := testDriverPaths(temp) + sec := secrets.New(p.SecretsFile) + fakeAS := &fakeDriverAgentserver{ + reg: agentserver.AgentRegistration{ + SandboxID: "sb-1", + TunnelToken: "tunnel-token", + ProxyToken: "sandbox-proxy-token", + WorkspaceID: "ws-reg", + ShortID: "abc123", + }, + whoamiErr: errors.New("whoami unavailable"), + } + + err := InstallDriver(context.Background(), DriverOptions{ + Paths: p, + Package: testDriverPackage(temp), + Secrets: sec, + ComputerName: "host", + AS: fakeAS, + ASOAuth: oauth.Config{Endpoint: "https://agent.test"}, + RequestDeviceCode: func(context.Context, oauth.Config) (oauth.DeviceCodeChallenge, error) { + return driverChallenge(), nil + }, + PollToken: func(context.Context, oauth.Config, oauth.DeviceCodeChallenge) (oauth.Token, error) { + return oauth.Token{AccessToken: workspaceToken("ws-claim", "Claim Workspace")}, nil + }, + Stdout: ioDiscard{}, + QR: driverMarkerQR("fake QR"), + }) + if err != nil { + t.Fatalf("InstallDriver: %v", err) + } + + st, err := state.NewStore(p.StateFile).Load() + if err != nil { + t.Fatal(err) + } + if st.Agentserver.WorkspaceID != "ws-claim" || st.Agentserver.WorkspaceName != "Claim Workspace" { + t.Fatalf("workspace state=%+v", st.Agentserver) + } +} + +type fakeDriverAgentserver struct { + reg agentserver.AgentRegistration + whoami agentserver.Identity + whoamiErr error + registerName string + registerType string +} + +func (f *fakeDriverAgentserver) RegisterAgent(_ context.Context, _ string, name, typ string) (agentserver.AgentRegistration, error) { + f.registerName = name + f.registerType = typ + return f.reg, nil +} + +func (f *fakeDriverAgentserver) Whoami(context.Context, string) (agentserver.Identity, error) { + if f.whoamiErr != nil { + return agentserver.Identity{}, f.whoamiErr + } + return f.whoami, nil +} + +func fakeRegisteredDriverAS() *fakeDriverAgentserver { + return &fakeDriverAgentserver{ + reg: agentserver.AgentRegistration{ + SandboxID: "sb-new", + TunnelToken: "tunnel-new", + ProxyToken: "proxy-new", + WorkspaceID: "ws-new", + ShortID: "new123", + }, + whoami: agentserver.Identity{ + Workspace: agentserver.Workspace{ID: "ws-new", Name: "New Workspace"}, + }, + } +} + +func testDriverPaths(root string) paths.Paths { + home := filepath.Join(root, "home") + installRoot := filepath.Join(home, ".agentserver-app") + return paths.Paths{ + UserHome: home, + InstallRoot: installRoot, + StateFile: filepath.Join(installRoot, "state.json"), + SecretsFile: filepath.Join(installRoot, "secrets.json"), + CodexConfigFile: filepath.Join(home, ".codex", "config.toml"), + } +} + +func testDriverPackage(root string) Package { + packageDir := filepath.Join(root, "pkg") + return Package{ + AgentserverExe: filepath.Join(packageDir, exeName("agentserver")), + PackageDir: packageDir, + DriverAgent: filepath.Join(packageDir, exeName("driver-agent")), + } +} + +func writeDriverState(t *testing.T, p paths.Paths, as state.AgentserverState) { + t.Helper() + if err := state.NewStore(p.StateFile).Update(func(s *state.State) error { + s.Agentserver = as + return nil + }); err != nil { + t.Fatal(err) + } +} + +func readTextFile(t *testing.T, path string) string { + t.Helper() + b, err := os.ReadFile(path) + if err != nil { + t.Fatal(err) + } + return string(b) +} + +func driverChallenge() oauth.DeviceCodeChallenge { + return oauth.DeviceCodeChallenge{ + DeviceCode: "device", + UserCode: "AS-CODE", + VerificationURIComplete: "https://agent.test/device?user_code=AS-CODE", + } +} + +func driverPollToken(context.Context, oauth.Config, oauth.DeviceCodeChallenge) (oauth.Token, error) { + return oauth.Token{AccessToken: "oauth-token"}, nil +} + +func driverMarkerQR(marker string) func(interface{ Write([]byte) (int, error) }, string) { + return func(w interface{ Write([]byte) (int, error) }, _ string) { + _, _ = w.Write([]byte(marker + "\n")) + } +} + +func workspaceToken(id, name string) string { + payload, err := json.Marshal(map[string]string{ + "workspace_id": id, + "workspace_name": name, + }) + if err != nil { + panic(err) + } + return "x." + base64.RawURLEncoding.EncodeToString(payload) + ".y" +} From 6be26ab62f69c650aef1c7f7b7623f277e4f6331 Mon Sep 17 00:00:00 2001 From: Zishu Yu Date: Sun, 14 Jun 2026 02:12:29 +0800 Subject: [PATCH 26/40] fix(headless): repair existing driver state --- internal/headless/driver.go | 34 ++++++++++++++++++ internal/headless/driver_test.go | 62 +++++++++++++++++++++++++++++++- 2 files changed, 95 insertions(+), 1 deletion(-) diff --git a/internal/headless/driver.go b/internal/headless/driver.go index 8dfde8d..9e212c4 100644 --- a/internal/headless/driver.go +++ b/internal/headless/driver.go @@ -105,6 +105,8 @@ func ensureDriver(ctx context.Context, opts DriverOptions, forceLogin bool) erro if err := registerDriver(ctx, opts); err != nil { return err } + } else if err := repairRegisteredDriverState(ctx, opts); err != nil { + return err } return refreshDriverFiles(opts.Paths, opts.Package, opts.Secrets) } @@ -213,6 +215,38 @@ func resolveDriverWorkspace(ctx context.Context, as AgentRegistrar, reg agentser return agentserver.Workspace{ID: reg.WorkspaceID} } +func repairRegisteredDriverState(ctx context.Context, opts DriverOptions) error { + proxyToken, _, err := driverSecrets(opts.Secrets) + if err != nil { + return err + } + var workspace agentserver.Workspace + if opts.AS != nil { + if identity, err := opts.AS.Whoami(ctx, proxyToken); err == nil { + workspace = identity.Workspace + } + } + return state.NewStore(opts.Paths.StateFile).Update(func(s *state.State) error { + if s.Agentserver.BaseURL == "" { + s.Agentserver.BaseURL = defaultASBaseURL(opts.ASOAuth) + } + if s.Agentserver.ShortID == "" { + s.Agentserver.ShortID = s.Agentserver.SandboxID + } + if s.Agentserver.WorkspaceID == "" && workspace.ID != "" { + s.Agentserver.WorkspaceID = workspace.ID + } + if s.Agentserver.WorkspaceName == "" && workspace.Name != "" && + (workspace.ID == "" || s.Agentserver.WorkspaceID == "" || workspace.ID == s.Agentserver.WorkspaceID) { + s.Agentserver.WorkspaceName = workspace.Name + } + if s.Agentserver.WorkspaceAPIKeySuffix == "" { + s.Agentserver.WorkspaceAPIKeySuffix = lastN(proxyToken, 4) + } + return nil + }) +} + func driverRegistered(p paths.Paths, sec secrets.Store) (bool, error) { st, err := state.NewStore(p.StateFile).Load() if err != nil { diff --git a/internal/headless/driver_test.go b/internal/headless/driver_test.go index f58789b..850ef6b 100644 --- a/internal/headless/driver_test.go +++ b/internal/headless/driver_test.go @@ -265,6 +265,64 @@ func TestInstallDriverSkipsDeviceFlowWhenAlreadyRegistered(t *testing.T) { } } +func TestInstallDriverRepairsExistingDriverState(t *testing.T) { + temp := t.TempDir() + p := testDriverPaths(temp) + sec := secrets.New(p.SecretsFile) + if err := sec.Set("agentserver_ws_api_key", "proxy-token"); err != nil { + t.Fatal(err) + } + if err := sec.Set("agentserver_tunnel_token", "tunnel-token"); err != nil { + t.Fatal(err) + } + writeDriverState(t, p, state.AgentserverState{ + SandboxID: "sb-1", + WorkspaceID: "ws-1", + }) + fakeAS := &fakeDriverAgentserver{ + whoami: agentserver.Identity{ + Workspace: agentserver.Workspace{ID: "ws-1", Name: "Workspace One"}, + }, + } + + err := InstallDriver(context.Background(), DriverOptions{ + Paths: p, + Package: testDriverPackage(temp), + Secrets: sec, + AS: fakeAS, + ASOAuth: oauth.Config{Endpoint: "https://agent.test"}, + RequestDeviceCode: func(context.Context, oauth.Config) (oauth.DeviceCodeChallenge, error) { + t.Fatal("RequestDeviceCode called for existing registration") + return oauth.DeviceCodeChallenge{}, nil + }, + PollToken: driverPollToken, + Stdout: ioDiscard{}, + }) + if err != nil { + t.Fatalf("InstallDriver: %v", err) + } + + st, err := state.NewStore(p.StateFile).Load() + if err != nil { + t.Fatal(err) + } + if st.Agentserver.BaseURL != "https://agent.test" { + t.Fatalf("BaseURL=%q, want https://agent.test", st.Agentserver.BaseURL) + } + if st.Agentserver.ShortID != "sb-1" { + t.Fatalf("ShortID=%q, want sandbox fallback", st.Agentserver.ShortID) + } + if st.Agentserver.WorkspaceName != "Workspace One" { + t.Fatalf("WorkspaceName=%q, want Workspace One", st.Agentserver.WorkspaceName) + } + if st.Agentserver.WorkspaceAPIKeySuffix != "oken" { + t.Fatalf("WorkspaceAPIKeySuffix=%q, want oken", st.Agentserver.WorkspaceAPIKeySuffix) + } + if fakeAS.whoamiToken != "proxy-token" { + t.Fatalf("Whoami token=%q, want proxy-token", fakeAS.whoamiToken) + } +} + func TestInstallDriverErrorsIfSecretsNil(t *testing.T) { err := InstallDriver(context.Background(), DriverOptions{ Paths: testDriverPaths(t.TempDir()), @@ -326,6 +384,7 @@ type fakeDriverAgentserver struct { reg agentserver.AgentRegistration whoami agentserver.Identity whoamiErr error + whoamiToken string registerName string registerType string } @@ -336,7 +395,8 @@ func (f *fakeDriverAgentserver) RegisterAgent(_ context.Context, _ string, name, return f.reg, nil } -func (f *fakeDriverAgentserver) Whoami(context.Context, string) (agentserver.Identity, error) { +func (f *fakeDriverAgentserver) Whoami(_ context.Context, token string) (agentserver.Identity, error) { + f.whoamiToken = token if f.whoamiErr != nil { return agentserver.Identity{}, f.whoamiErr } From 18274b7974ab9ae8c5f5d3ee13ce3e94b3e2ca00 Mon Sep 17 00:00:00 2001 From: Zishu Yu Date: Sun, 14 Jun 2026 02:16:08 +0800 Subject: [PATCH 27/40] fix(headless): repair driver workspace metadata --- internal/headless/driver.go | 2 +- internal/headless/driver_test.go | 47 ++++++++++++++++++++++++++++++++ 2 files changed, 48 insertions(+), 1 deletion(-) diff --git a/internal/headless/driver.go b/internal/headless/driver.go index 9e212c4..4e9a15d 100644 --- a/internal/headless/driver.go +++ b/internal/headless/driver.go @@ -252,7 +252,7 @@ func driverRegistered(p paths.Paths, sec secrets.Store) (bool, error) { if err != nil { return false, err } - if strings.TrimSpace(st.Agentserver.SandboxID) == "" || strings.TrimSpace(st.Agentserver.WorkspaceID) == "" { + if strings.TrimSpace(st.Agentserver.SandboxID) == "" { return false, nil } if ok, err := secretPresent(sec, driverProxySecretKey); err != nil || !ok { diff --git a/internal/headless/driver_test.go b/internal/headless/driver_test.go index 850ef6b..b26e206 100644 --- a/internal/headless/driver_test.go +++ b/internal/headless/driver_test.go @@ -323,6 +323,53 @@ func TestInstallDriverRepairsExistingDriverState(t *testing.T) { } } +func TestInstallDriverRepairsMissingWorkspaceIDWithoutDeviceFlow(t *testing.T) { + temp := t.TempDir() + p := testDriverPaths(temp) + sec := secrets.New(p.SecretsFile) + if err := sec.Set("agentserver_ws_api_key", "proxy-token"); err != nil { + t.Fatal(err) + } + if err := sec.Set("agentserver_tunnel_token", "tunnel-token"); err != nil { + t.Fatal(err) + } + writeDriverState(t, p, state.AgentserverState{ + BaseURL: "https://agent.test", + SandboxID: "sb-1", + ShortID: "abc123", + }) + fakeAS := &fakeDriverAgentserver{ + whoami: agentserver.Identity{ + Workspace: agentserver.Workspace{ID: "ws-1", Name: "Workspace One"}, + }, + } + + err := InstallDriver(context.Background(), DriverOptions{ + Paths: p, + Package: testDriverPackage(temp), + Secrets: sec, + AS: fakeAS, + ASOAuth: oauth.Config{Endpoint: "https://agent.test"}, + RequestDeviceCode: func(context.Context, oauth.Config) (oauth.DeviceCodeChallenge, error) { + t.Fatal("RequestDeviceCode called for repairable existing registration") + return oauth.DeviceCodeChallenge{}, nil + }, + PollToken: driverPollToken, + Stdout: ioDiscard{}, + }) + if err != nil { + t.Fatalf("InstallDriver: %v", err) + } + + st, err := state.NewStore(p.StateFile).Load() + if err != nil { + t.Fatal(err) + } + if st.Agentserver.WorkspaceID != "ws-1" || st.Agentserver.WorkspaceName != "Workspace One" { + t.Fatalf("workspace state=%+v", st.Agentserver) + } +} + func TestInstallDriverErrorsIfSecretsNil(t *testing.T) { err := InstallDriver(context.Background(), DriverOptions{ Paths: testDriverPaths(t.TempDir()), From afe3230dadef2f727538d12654f1bb6e62cb1d10 Mon Sep 17 00:00:00 2001 From: Zishu Yu Date: Sun, 14 Jun 2026 02:24:54 +0800 Subject: [PATCH 28/40] fix(headless): harden driver install paths --- internal/headless/driver.go | 13 ++- internal/headless/driver_test.go | 191 +++++++++++++++++++++++++++++++ 2 files changed, 201 insertions(+), 3 deletions(-) diff --git a/internal/headless/driver.go b/internal/headless/driver.go index 4e9a15d..414ba04 100644 --- a/internal/headless/driver.go +++ b/internal/headless/driver.go @@ -97,11 +97,18 @@ func ensureDriver(ctx context.Context, opts DriverOptions, forceLogin bool) erro } opts = defaultDriverOptions(opts) + if forceLogin { + if err := registerDriver(ctx, opts); err != nil { + return err + } + return refreshDriverFiles(opts.Paths, opts.Package, opts.Secrets) + } + registered, err := driverRegistered(opts.Paths, opts.Secrets) if err != nil { return err } - if forceLogin || !registered { + if !registered { if err := registerDriver(ctx, opts); err != nil { return err } @@ -178,7 +185,7 @@ func registerDriver(ctx context.Context, opts DriverOptions) error { if err != nil { return fmt.Errorf("register agentserver agent: %w", err) } - if reg.ProxyToken == "" || reg.TunnelToken == "" || reg.SandboxID == "" || reg.ShortID == "" { + if reg.ProxyToken == "" || reg.TunnelToken == "" || reg.SandboxID == "" { return errors.New("register agentserver agent: incomplete registration") } @@ -195,7 +202,7 @@ func registerDriver(ctx context.Context, opts DriverOptions) error { return state.NewStore(opts.Paths.StateFile).Update(func(s *state.State) error { s.Agentserver.BaseURL = defaultASBaseURL(opts.ASOAuth) s.Agentserver.SandboxID = reg.SandboxID - s.Agentserver.ShortID = reg.ShortID + s.Agentserver.ShortID = defaultString(reg.ShortID, reg.SandboxID) s.Agentserver.WorkspaceID = workspace.ID s.Agentserver.WorkspaceName = workspace.Name s.Agentserver.WorkspaceAPIKeySuffix = lastN(reg.ProxyToken, 4) diff --git a/internal/headless/driver_test.go b/internal/headless/driver_test.go index b26e206..dc673ab 100644 --- a/internal/headless/driver_test.go +++ b/internal/headless/driver_test.go @@ -1,11 +1,14 @@ package headless import ( + "archive/tar" "bytes" + "compress/gzip" "context" "encoding/base64" "encoding/json" "errors" + "io" "os" "path/filepath" "strings" @@ -265,6 +268,60 @@ func TestInstallDriverSkipsDeviceFlowWhenAlreadyRegistered(t *testing.T) { } } +func TestInstallDriverWritesPersistentConfigAndSupport(t *testing.T) { + temp := t.TempDir() + p := testDriverPaths(temp) + sec := secrets.New(p.SecretsFile) + pkg := testDriverPackage(temp) + writeTestTarGz(t, filepath.Join(pkg.PackageDir, "driver-skills.tar.gz"), map[string]string{ + "skills/demo/SKILL.md": "demo skill", + }) + writeTestTarGz(t, filepath.Join(pkg.PackageDir, "driver-superpower-skills.tar.gz"), map[string]string{ + "skills/super/SKILL.md": "super skill", + }) + writeTestTarGz(t, filepath.Join(pkg.PackageDir, "driver-codex-prompts.tar.gz"), map[string]string{ + "prompts-codex/AGENTS.md": "driver prompt", + }) + + err := InstallDriver(context.Background(), DriverOptions{ + Paths: p, + Package: pkg, + Secrets: sec, + ComputerName: "host", + AS: fakeRegisteredDriverAS(), + ASOAuth: oauth.Config{Endpoint: "https://agent.test"}, + RequestDeviceCode: func(context.Context, oauth.Config) (oauth.DeviceCodeChallenge, error) { + return driverChallenge(), nil + }, + PollToken: driverPollToken, + Stdout: ioDiscard{}, + QR: driverMarkerQR("fake QR"), + }) + if err != nil { + t.Fatalf("InstallDriver: %v", err) + } + + driverConfig := readTextFile(t, filepath.Join(p.UserHome, ".config", "multi-agent", "driver.yaml")) + for _, want := range []string{ + `url: "https://agent.test"`, + `proxy_token: "proxy-new"`, + `workdir: "` + filepath.ToSlash(p.UserHome) + `"`, + } { + if !strings.Contains(driverConfig, want) { + t.Fatalf("persistent driver config missing %q:\n%s", want, driverConfig) + } + } + for _, path := range []string{ + filepath.Join(p.UserHome, ".agents", "skills", "demo", "SKILL.md"), + filepath.Join(p.UserHome, ".codex", "skills", "super", "SKILL.md"), + filepath.Join(p.UserHome, ".codex", "AGENTS.md"), + } { + if _, err := os.Stat(path); err != nil { + t.Fatalf("expected installed support file %s: %v", path, err) + } + } +} + func TestInstallDriverRepairsExistingDriverState(t *testing.T) { temp := t.TempDir() p := testDriverPaths(temp) @@ -370,6 +427,47 @@ func TestInstallDriverRepairsMissingWorkspaceIDWithoutDeviceFlow(t *testing.T) { } } +func TestSwitchWorkspaceBypassesBrokenExistingSecretReads(t *testing.T) { + temp := t.TempDir() + p := testDriverPaths(temp) + sec := &writeableBrokenReadSecrets{ + err: errors.New("old secrets unreadable"), + values: map[string]string{}, + } + writeDriverState(t, p, state.AgentserverState{ + BaseURL: "https://agent.old", + SandboxID: "sb-old", + ShortID: "old123", + WorkspaceID: "ws-old", + }) + calls := 0 + + err := SwitchWorkspace(context.Background(), DriverOptions{ + Paths: p, + Package: testDriverPackage(temp), + Secrets: sec, + ComputerName: "host", + AS: fakeRegisteredDriverAS(), + ASOAuth: oauth.Config{Endpoint: "https://agent.test"}, + RequestDeviceCode: func(context.Context, oauth.Config) (oauth.DeviceCodeChallenge, error) { + calls++ + return driverChallenge(), nil + }, + PollToken: driverPollToken, + Stdout: ioDiscard{}, + QR: driverMarkerQR("fake QR"), + }) + if err != nil { + t.Fatalf("SwitchWorkspace: %v", err) + } + if calls != 1 { + t.Fatalf("RequestDeviceCode calls=%d, want 1", calls) + } + if got := sec.values["agentserver_ws_api_key"]; got != "proxy-new" { + t.Fatalf("stored proxy token=%q, want proxy-new", got) + } +} + func TestInstallDriverErrorsIfSecretsNil(t *testing.T) { err := InstallDriver(context.Background(), DriverOptions{ Paths: testDriverPaths(t.TempDir()), @@ -427,6 +525,49 @@ func TestInstallDriverStoresWorkspaceFallbackFromAccessTokenClaim(t *testing.T) } } +func TestInstallDriverUsesSandboxIDWhenRegistrationOmitsShortID(t *testing.T) { + temp := t.TempDir() + p := testDriverPaths(temp) + sec := secrets.New(p.SecretsFile) + fakeAS := &fakeDriverAgentserver{ + reg: agentserver.AgentRegistration{ + SandboxID: "sb-1", + TunnelToken: "tunnel-token", + ProxyToken: "sandbox-proxy-token", + WorkspaceID: "ws-reg", + }, + whoami: agentserver.Identity{ + Workspace: agentserver.Workspace{ID: "ws-1", Name: "Workspace One"}, + }, + } + + err := InstallDriver(context.Background(), DriverOptions{ + Paths: p, + Package: testDriverPackage(temp), + Secrets: sec, + ComputerName: "host", + AS: fakeAS, + ASOAuth: oauth.Config{Endpoint: "https://agent.test"}, + RequestDeviceCode: func(context.Context, oauth.Config) (oauth.DeviceCodeChallenge, error) { + return driverChallenge(), nil + }, + PollToken: driverPollToken, + Stdout: ioDiscard{}, + QR: driverMarkerQR("fake QR"), + }) + if err != nil { + t.Fatalf("InstallDriver: %v", err) + } + + st, err := state.NewStore(p.StateFile).Load() + if err != nil { + t.Fatal(err) + } + if st.Agentserver.ShortID != "sb-1" { + t.Fatalf("ShortID=%q, want sandbox fallback", st.Agentserver.ShortID) + } +} + type fakeDriverAgentserver struct { reg agentserver.AgentRegistration whoami agentserver.Identity @@ -533,3 +674,53 @@ func workspaceToken(id, name string) string { } return "x." + base64.RawURLEncoding.EncodeToString(payload) + ".y" } + +type writeableBrokenReadSecrets struct { + err error + values map[string]string +} + +func (s *writeableBrokenReadSecrets) Get(key string) (string, error) { + if v, ok := s.values[key]; ok { + return v, nil + } + return "", s.err +} + +func (s *writeableBrokenReadSecrets) Set(key, value string) error { + s.values[key] = value + return nil +} + +func (s *writeableBrokenReadSecrets) Delete(key string) error { + delete(s.values, key) + return nil +} + +func writeTestTarGz(t *testing.T, path string, files map[string]string) { + t.Helper() + if err := os.MkdirAll(filepath.Dir(path), 0o755); err != nil { + t.Fatal(err) + } + f, err := os.Create(path) + if err != nil { + t.Fatal(err) + } + defer f.Close() + gz := gzip.NewWriter(f) + defer gz.Close() + tw := tar.NewWriter(gz) + defer tw.Close() + for name, body := range files { + if err := tw.WriteHeader(&tar.Header{ + Name: name, + Mode: 0o644, + Size: int64(len(body)), + }); err != nil { + t.Fatal(err) + } + if _, err := io.WriteString(tw, body); err != nil { + t.Fatal(err) + } + } +} From 184cab81904b03cace600dc8032f5c85d4016539 Mon Sep 17 00:00:00 2001 From: Zishu Yu Date: Sun, 14 Jun 2026 02:33:04 +0800 Subject: [PATCH 29/40] fix(headless): refresh driver workspace state --- internal/headless/driver.go | 5 ++- internal/headless/driver_test.go | 53 ++++++++++++++++++++++++++++++++ 2 files changed, 55 insertions(+), 3 deletions(-) diff --git a/internal/headless/driver.go b/internal/headless/driver.go index 414ba04..ea38875 100644 --- a/internal/headless/driver.go +++ b/internal/headless/driver.go @@ -240,11 +240,10 @@ func repairRegisteredDriverState(ctx context.Context, opts DriverOptions) error if s.Agentserver.ShortID == "" { s.Agentserver.ShortID = s.Agentserver.SandboxID } - if s.Agentserver.WorkspaceID == "" && workspace.ID != "" { + if workspace.ID != "" { s.Agentserver.WorkspaceID = workspace.ID } - if s.Agentserver.WorkspaceName == "" && workspace.Name != "" && - (workspace.ID == "" || s.Agentserver.WorkspaceID == "" || workspace.ID == s.Agentserver.WorkspaceID) { + if workspace.Name != "" && workspace.ID == s.Agentserver.WorkspaceID { s.Agentserver.WorkspaceName = workspace.Name } if s.Agentserver.WorkspaceAPIKeySuffix == "" { diff --git a/internal/headless/driver_test.go b/internal/headless/driver_test.go index dc673ab..a0199a2 100644 --- a/internal/headless/driver_test.go +++ b/internal/headless/driver_test.go @@ -427,6 +427,59 @@ func TestInstallDriverRepairsMissingWorkspaceIDWithoutDeviceFlow(t *testing.T) { } } +func TestInstallDriverRepairsStaleWorkspaceFromCurrentSecrets(t *testing.T) { + temp := t.TempDir() + p := testDriverPaths(temp) + sec := secrets.New(p.SecretsFile) + if err := sec.Set("agentserver_ws_api_key", "proxy-new"); err != nil { + t.Fatal(err) + } + if err := sec.Set("agentserver_tunnel_token", "tunnel-new"); err != nil { + t.Fatal(err) + } + writeDriverState(t, p, state.AgentserverState{ + BaseURL: "https://agent.test", + SandboxID: "sb-1", + ShortID: "abc123", + WorkspaceID: "ws-old", + WorkspaceName: "Old Workspace", + }) + fakeAS := &fakeDriverAgentserver{ + whoami: agentserver.Identity{ + Workspace: agentserver.Workspace{ID: "ws-new", Name: "New Workspace"}, + }, + } + + err := InstallDriver(context.Background(), DriverOptions{ + Paths: p, + Package: testDriverPackage(temp), + Secrets: sec, + AS: fakeAS, + ASOAuth: oauth.Config{Endpoint: "https://agent.test"}, + RequestDeviceCode: func(context.Context, oauth.Config) (oauth.DeviceCodeChallenge, error) { + t.Fatal("RequestDeviceCode called for repairable existing registration") + return oauth.DeviceCodeChallenge{}, nil + }, + PollToken: driverPollToken, + Stdout: ioDiscard{}, + }) + if err != nil { + t.Fatalf("InstallDriver: %v", err) + } + + st, err := state.NewStore(p.StateFile).Load() + if err != nil { + t.Fatal(err) + } + if st.Agentserver.WorkspaceID != "ws-new" || st.Agentserver.WorkspaceName != "New Workspace" { + t.Fatalf("workspace state=%+v", st.Agentserver) + } + driverConfig := readTextFile(t, filepath.Join(p.UserHome, ".config", "multi-agent", "driver.yaml")) + if !strings.Contains(driverConfig, `workspace_id: "ws-new"`) { + t.Fatalf("driver config kept stale workspace:\n%s", driverConfig) + } +} + func TestSwitchWorkspaceBypassesBrokenExistingSecretReads(t *testing.T) { temp := t.TempDir() p := testDriverPaths(temp) From f29e378de8ae25c0afe259de633a91590903fef2 Mon Sep 17 00:00:00 2001 From: Zishu Yu Date: Sun, 14 Jun 2026 02:36:31 +0800 Subject: [PATCH 30/40] fix(headless): refresh driver repair metadata --- internal/headless/driver.go | 8 ++--- internal/headless/driver_test.go | 61 +++++++++++++++++++++++++++++--- 2 files changed, 58 insertions(+), 11 deletions(-) diff --git a/internal/headless/driver.go b/internal/headless/driver.go index ea38875..9f0af74 100644 --- a/internal/headless/driver.go +++ b/internal/headless/driver.go @@ -234,9 +234,7 @@ func repairRegisteredDriverState(ctx context.Context, opts DriverOptions) error } } return state.NewStore(opts.Paths.StateFile).Update(func(s *state.State) error { - if s.Agentserver.BaseURL == "" { - s.Agentserver.BaseURL = defaultASBaseURL(opts.ASOAuth) - } + s.Agentserver.BaseURL = defaultASBaseURL(opts.ASOAuth) if s.Agentserver.ShortID == "" { s.Agentserver.ShortID = s.Agentserver.SandboxID } @@ -246,9 +244,7 @@ func repairRegisteredDriverState(ctx context.Context, opts DriverOptions) error if workspace.Name != "" && workspace.ID == s.Agentserver.WorkspaceID { s.Agentserver.WorkspaceName = workspace.Name } - if s.Agentserver.WorkspaceAPIKeySuffix == "" { - s.Agentserver.WorkspaceAPIKeySuffix = lastN(proxyToken, 4) - } + s.Agentserver.WorkspaceAPIKeySuffix = lastN(proxyToken, 4) return nil }) } diff --git a/internal/headless/driver_test.go b/internal/headless/driver_test.go index a0199a2..ef33403 100644 --- a/internal/headless/driver_test.go +++ b/internal/headless/driver_test.go @@ -438,11 +438,12 @@ func TestInstallDriverRepairsStaleWorkspaceFromCurrentSecrets(t *testing.T) { t.Fatal(err) } writeDriverState(t, p, state.AgentserverState{ - BaseURL: "https://agent.test", - SandboxID: "sb-1", - ShortID: "abc123", - WorkspaceID: "ws-old", - WorkspaceName: "Old Workspace", + BaseURL: "https://agent.test", + SandboxID: "sb-1", + ShortID: "abc123", + WorkspaceID: "ws-old", + WorkspaceName: "Old Workspace", + WorkspaceAPIKeySuffix: "old!", }) fakeAS := &fakeDriverAgentserver{ whoami: agentserver.Identity{ @@ -474,12 +475,62 @@ func TestInstallDriverRepairsStaleWorkspaceFromCurrentSecrets(t *testing.T) { if st.Agentserver.WorkspaceID != "ws-new" || st.Agentserver.WorkspaceName != "New Workspace" { t.Fatalf("workspace state=%+v", st.Agentserver) } + if st.Agentserver.WorkspaceAPIKeySuffix != "-new" { + t.Fatalf("WorkspaceAPIKeySuffix=%q, want -new", st.Agentserver.WorkspaceAPIKeySuffix) + } driverConfig := readTextFile(t, filepath.Join(p.UserHome, ".config", "multi-agent", "driver.yaml")) if !strings.Contains(driverConfig, `workspace_id: "ws-new"`) { t.Fatalf("driver config kept stale workspace:\n%s", driverConfig) } } +func TestInstallDriverRepairsStaleBaseURL(t *testing.T) { + temp := t.TempDir() + p := testDriverPaths(temp) + sec := secrets.New(p.SecretsFile) + if err := sec.Set("agentserver_ws_api_key", "proxy-token"); err != nil { + t.Fatal(err) + } + if err := sec.Set("agentserver_tunnel_token", "tunnel-token"); err != nil { + t.Fatal(err) + } + writeDriverState(t, p, state.AgentserverState{ + BaseURL: "https://agent.old", + SandboxID: "sb-1", + ShortID: "abc123", + WorkspaceID: "ws-1", + }) + + err := InstallDriver(context.Background(), DriverOptions{ + Paths: p, + Package: testDriverPackage(temp), + Secrets: sec, + AS: fakeRegisteredDriverAS(), + ASOAuth: oauth.Config{Endpoint: "https://agent.test"}, + RequestDeviceCode: func(context.Context, oauth.Config) (oauth.DeviceCodeChallenge, error) { + t.Fatal("RequestDeviceCode called for repairable existing registration") + return oauth.DeviceCodeChallenge{}, nil + }, + PollToken: driverPollToken, + Stdout: ioDiscard{}, + }) + if err != nil { + t.Fatalf("InstallDriver: %v", err) + } + + st, err := state.NewStore(p.StateFile).Load() + if err != nil { + t.Fatal(err) + } + if st.Agentserver.BaseURL != "https://agent.test" { + t.Fatalf("BaseURL=%q, want https://agent.test", st.Agentserver.BaseURL) + } + driverConfig := readTextFile(t, filepath.Join(p.UserHome, ".config", "multi-agent", "driver.yaml")) + if strings.Contains(driverConfig, "https://agent.old") || !strings.Contains(driverConfig, `url: "https://agent.test"`) { + t.Fatalf("driver config kept stale base URL:\n%s", driverConfig) + } +} + func TestSwitchWorkspaceBypassesBrokenExistingSecretReads(t *testing.T) { temp := t.TempDir() p := testDriverPaths(temp) From c1e652f774b5f6f487737f55adbe7e4dde7ef72b Mon Sep 17 00:00:00 2001 From: Zishu Yu Date: Sun, 14 Jun 2026 02:48:09 +0800 Subject: [PATCH 31/40] feat(headless): add agentserver CLI --- cmd/agentserver/main.go | 177 +++++++++++++++++++++++++++++++++++ cmd/agentserver/main_test.go | 111 ++++++++++++++++++++++ 2 files changed, 288 insertions(+) create mode 100644 cmd/agentserver/main.go create mode 100644 cmd/agentserver/main_test.go diff --git a/cmd/agentserver/main.go b/cmd/agentserver/main.go new file mode 100644 index 0000000..b62410d --- /dev/null +++ b/cmd/agentserver/main.go @@ -0,0 +1,177 @@ +package main + +import ( + "bufio" + "context" + "errors" + "fmt" + "io" + "log" + "os" + "path/filepath" + "strings" + + "github.com/agentserver/agentserver-pkg/internal/headless" + "github.com/agentserver/agentserver-pkg/internal/modelaccess" + "github.com/agentserver/agentserver-pkg/internal/modelserver" + "github.com/agentserver/agentserver-pkg/internal/oauth" + "github.com/agentserver/agentserver-pkg/internal/paths" + "github.com/agentserver/agentserver-pkg/internal/secrets" + "github.com/agentserver/agentserver-pkg/internal/terminalauth" +) + +type app struct { + ensureAccess func(context.Context) error + runSlave func(context.Context) error + installDriver func(context.Context) error + switchWorkspace func(context.Context) error + serveDriverMCP func(context.Context) error + runDaemon func(context.Context) error +} + +func main() { + if err := newApp().run(context.Background(), os.Args[1:]); err != nil { + log.Fatalf("agentserver: %v", err) + } +} + +func newApp() app { + p, err := paths.Default() + if err != nil { + return app{ + ensureAccess: func(context.Context) error { return err }, + runDaemon: func(context.Context) error { return err }, + } + } + exe, err := os.Executable() + if err != nil { + return app{ + ensureAccess: func(context.Context) error { return err }, + runDaemon: func(context.Context) error { return err }, + } + } + pkg := headless.PackagePaths(exe) + sec := secrets.New(p.SecretsFile) + cachedCodex := "" + + return app{ + ensureAccess: func(ctx context.Context) error { + if _, err := modelaccess.Ensure(ctx, modelaccess.EnsureOptions{ + CodexConfigPath: p.CodexConfigFile, + Secrets: sec, + PrintChallenge: func(title string, ch oauth.DeviceCodeChallenge) { + terminalauth.PrintChallenge(os.Stdout, title, ch, terminalauth.DefaultQR) + }, + StartDaemon: func(ctx context.Context) error { + return modelaccess.EnsureDaemon(ctx, modelaccess.EnsureDaemonOptions{ + ExePath: pkg.AgentserverExe, + ProxyBaseURL: "http://127.0.0.1:53452", + }) + }, + }); err != nil { + return err + } + codexRuntime, err := headless.ResolveCodex(ctx, headless.CodexResolveOptions{ + Paths: p, + Package: pkg, + }) + if err != nil { + return err + } + cachedCodex = codexRuntime.Path + return nil + }, + runSlave: func(ctx context.Context) error { + wd, _ := os.Getwd() + return headless.RunSlave(ctx, headless.SlaveOptions{ + Paths: p, + Package: pkg, + WorkDir: wd, + NamePrompt: promptName, + Stdout: os.Stdout, + QR: terminalauth.DefaultQR, + CodexBin: cachedCodex, + }) + }, + installDriver: func(ctx context.Context) error { + return headless.InstallDriver(ctx, headless.DriverOptions{ + Paths: p, + Package: pkg, + Secrets: sec, + Stdout: os.Stdout, + QR: terminalauth.DefaultQR, + }) + }, + switchWorkspace: func(ctx context.Context) error { + return headless.SwitchWorkspace(ctx, headless.DriverOptions{ + Paths: p, + Package: pkg, + Secrets: sec, + Stdout: os.Stdout, + QR: terminalauth.DefaultQR, + }) + }, + serveDriverMCP: func(ctx context.Context) error { + wd, _ := os.Getwd() + return headless.ServeDriverMCP(ctx, headless.DriverMCPOptions{ + Paths: p, + Package: pkg, + Secrets: sec, + WorkDir: wd, + }) + }, + runDaemon: func(ctx context.Context) error { + return modelaccess.RunDaemon(ctx, modelaccess.DaemonOptions{ + Secrets: sec, + OAuth: modelserver.OAuthConfig(), + LockPath: filepath.Join(p.InstallRoot, "token-refresher.lock"), + }) + }, + } +} + +func (a app) run(ctx context.Context, args []string) error { + cmd := "" + if len(args) > 0 { + cmd = args[0] + } + if cmd == "model-proxy-daemon" { + return a.runDaemon(ctx) + } + + if a.ensureAccess != nil { + if err := a.ensureAccess(ctx); err != nil { + return err + } + } + switch cmd { + case "": + return a.runSlave(ctx) + case "install-driver": + return a.installDriver(ctx) + case "switch-workspace": + return a.switchWorkspace(ctx) + case "serve-driver-mcp": + return a.serveDriverMCP(ctx) + default: + return fmt.Errorf("unknown command %q", cmd) + } +} + +func promptName(defaultName string) (string, error) { + fmt.Fprintf(os.Stdout, "Slave name [%s]: ", defaultName) + line, err := bufio.NewReader(os.Stdin).ReadString('\n') + if err != nil { + if !errors.Is(err, io.EOF) { + return "", err + } + if line == "" { + return defaultName, nil + } + } + name := strings.TrimSpace(line) + if name == "" { + return defaultName, nil + } + return name, nil +} diff --git a/cmd/agentserver/main_test.go b/cmd/agentserver/main_test.go new file mode 100644 index 0000000..7eca726 --- /dev/null +++ b/cmd/agentserver/main_test.go @@ -0,0 +1,111 @@ +package main + +import ( + "context" + "strings" + "testing" +) + +func TestRunCommandEnsuresAccessForUserCommands(t *testing.T) { + for _, tc := range []struct { + name string + args []string + command func(*app, *int, *int, *testing.T) + }{ + { + name: "default", + args: []string{}, + command: func(app *app, access, calls *int, t *testing.T) { + app.runSlave = func(context.Context) error { + if *access != 1 { + t.Fatalf("runSlave called after ensureAccess calls=%d, want 1", *access) + } + *calls = *calls + 1 + return nil + } + }, + }, + { + name: "install-driver", + args: []string{"install-driver"}, + command: func(app *app, access, calls *int, t *testing.T) { + app.installDriver = func(context.Context) error { + if *access != 1 { + t.Fatalf("installDriver called after ensureAccess calls=%d, want 1", *access) + } + *calls = *calls + 1 + return nil + } + }, + }, + { + name: "switch-workspace", + args: []string{"switch-workspace"}, + command: func(app *app, access, calls *int, t *testing.T) { + app.switchWorkspace = func(context.Context) error { + if *access != 1 { + t.Fatalf("switchWorkspace called after ensureAccess calls=%d, want 1", *access) + } + *calls = *calls + 1 + return nil + } + }, + }, + { + name: "serve-driver-mcp", + args: []string{"serve-driver-mcp"}, + command: func(app *app, access, calls *int, t *testing.T) { + app.serveDriverMCP = func(context.Context) error { + if *access != 1 { + t.Fatalf("serveDriverMCP called after ensureAccess calls=%d, want 1", *access) + } + *calls = *calls + 1 + return nil + } + }, + }, + } { + t.Run(strings.ReplaceAll(tc.name, "-", "_"), func(t *testing.T) { + access := 0 + commandCalls := 0 + app := app{ + ensureAccess: func(context.Context) error { access++; return nil }, + } + tc.command(&app, &access, &commandCalls, t) + if err := app.run(context.Background(), tc.args); err != nil { + t.Fatalf("run: %v", err) + } + if access != 1 { + t.Fatalf("ensureAccess calls=%d, want 1", access) + } + if commandCalls != 1 { + t.Fatalf("command calls=%d, want 1", commandCalls) + } + }) + } +} + +func TestRunCommandSkipsAccessForDaemon(t *testing.T) { + access := 0 + daemonCalls := 0 + app := app{ + ensureAccess: func(context.Context) error { access++; return nil }, + runDaemon: func(context.Context) error { daemonCalls++; return nil }, + } + if err := app.run(context.Background(), []string{"model-proxy-daemon"}); err != nil { + t.Fatalf("run: %v", err) + } + if access != 0 { + t.Fatalf("ensureAccess calls=%d, want 0", access) + } + if daemonCalls != 1 { + t.Fatalf("runDaemon calls=%d, want 1", daemonCalls) + } +} + +func TestRunCommandRejectsUnknownCommand(t *testing.T) { + app := app{} + if err := app.run(context.Background(), []string{"unknown"}); err == nil { + t.Fatal("expected unknown command error") + } +} From 8061ada75ced21695024f42844b74cecbeaf85c1 Mon Sep 17 00:00:00 2001 From: Zishu Yu Date: Sun, 14 Jun 2026 02:58:14 +0800 Subject: [PATCH 32/40] fix(headless): harden agentserver CLI routing --- cmd/agentserver/main.go | 30 +++++++++++++-- cmd/agentserver/main_test.go | 72 +++++++++++++++++++++++++++++++++++- 2 files changed, 97 insertions(+), 5 deletions(-) diff --git a/cmd/agentserver/main.go b/cmd/agentserver/main.go index b62410d..286697e 100644 --- a/cmd/agentserver/main.go +++ b/cmd/agentserver/main.go @@ -22,6 +22,7 @@ import ( type app struct { ensureAccess func(context.Context) error + ensureCodex func(context.Context) error runSlave func(context.Context) error installDriver func(context.Context) error switchWorkspace func(context.Context) error @@ -71,6 +72,9 @@ func newApp() app { }); err != nil { return err } + return nil + }, + ensureCodex: func(ctx context.Context) error { codexRuntime, err := headless.ResolveCodex(ctx, headless.CodexResolveOptions{ Paths: p, Package: pkg, @@ -138,12 +142,22 @@ func (a app) run(ctx context.Context, args []string) error { if cmd == "model-proxy-daemon" { return a.runDaemon(ctx) } + switch cmd { + case "", "install-driver", "switch-workspace", "serve-driver-mcp": + default: + return fmt.Errorf("unknown command %q", cmd) + } if a.ensureAccess != nil { if err := a.ensureAccess(ctx); err != nil { return err } } + if cmd == "" && a.ensureCodex != nil { + if err := a.ensureCodex(ctx); err != nil { + return err + } + } switch cmd { case "": return a.runSlave(ctx) @@ -153,14 +167,22 @@ func (a app) run(ctx context.Context, args []string) error { return a.switchWorkspace(ctx) case "serve-driver-mcp": return a.serveDriverMCP(ctx) - default: - return fmt.Errorf("unknown command %q", cmd) } + return nil } func promptName(defaultName string) (string, error) { - fmt.Fprintf(os.Stdout, "Slave name [%s]: ", defaultName) - line, err := bufio.NewReader(os.Stdin).ReadString('\n') + return promptNameWithIO(os.Stdin, os.Stdout, defaultName) +} + +func promptNameWithIO(r io.Reader, w io.Writer, defaultName string) (string, error) { + if w != nil { + fmt.Fprintf(w, "Slave name [%s]: ", defaultName) + } + if r == nil { + return defaultName, nil + } + line, err := bufio.NewReader(r).ReadString('\n') if err != nil { if !errors.Is(err, io.EOF) { return "", err diff --git a/cmd/agentserver/main_test.go b/cmd/agentserver/main_test.go index 7eca726..6c586c1 100644 --- a/cmd/agentserver/main_test.go +++ b/cmd/agentserver/main_test.go @@ -1,6 +1,7 @@ package main import ( + "bytes" "context" "strings" "testing" @@ -85,6 +86,43 @@ func TestRunCommandEnsuresAccessForUserCommands(t *testing.T) { } } +func TestRunCommandResolvesCodexOnlyForDefaultSlave(t *testing.T) { + for _, tc := range []struct { + name string + args []string + wantCodex int + }{ + {name: "default", args: []string{}, wantCodex: 1}, + {name: "install-driver", args: []string{"install-driver"}}, + {name: "switch-workspace", args: []string{"switch-workspace"}}, + {name: "serve-driver-mcp", args: []string{"serve-driver-mcp"}}, + } { + t.Run(strings.ReplaceAll(tc.name, "-", "_"), func(t *testing.T) { + codex := 0 + app := app{ + ensureAccess: func(context.Context) error { return nil }, + ensureCodex: func(context.Context) error { codex++; return nil }, + runSlave: func(context.Context) error { return nil }, + installDriver: func(context.Context) error { + return nil + }, + switchWorkspace: func(context.Context) error { + return nil + }, + serveDriverMCP: func(context.Context) error { + return nil + }, + } + if err := app.run(context.Background(), tc.args); err != nil { + t.Fatalf("run: %v", err) + } + if codex != tc.wantCodex { + t.Fatalf("ensureCodex calls=%d, want %d", codex, tc.wantCodex) + } + }) + } +} + func TestRunCommandSkipsAccessForDaemon(t *testing.T) { access := 0 daemonCalls := 0 @@ -104,8 +142,40 @@ func TestRunCommandSkipsAccessForDaemon(t *testing.T) { } func TestRunCommandRejectsUnknownCommand(t *testing.T) { - app := app{} + access := 0 + app := app{ + ensureAccess: func(context.Context) error { access++; return nil }, + } if err := app.run(context.Background(), []string{"unknown"}); err == nil { t.Fatal("expected unknown command error") } + if access != 0 { + t.Fatalf("ensureAccess calls=%d, want 0", access) + } +} + +func TestPromptNameWithIO(t *testing.T) { + for _, tc := range []struct { + name string + input string + want string + }{ + {name: "eof", input: "", want: "default"}, + {name: "blank", input: "\n", want: "default"}, + {name: "trimmed", input: " custom \n", want: "custom"}, + } { + t.Run(tc.name, func(t *testing.T) { + var out bytes.Buffer + got, err := promptNameWithIO(strings.NewReader(tc.input), &out, "default") + if err != nil { + t.Fatalf("promptNameWithIO: %v", err) + } + if got != tc.want { + t.Fatalf("name=%q, want %q", got, tc.want) + } + if !strings.Contains(out.String(), "Slave name [default]: ") { + t.Fatalf("prompt output=%q", out.String()) + } + }) + } } From b09bb4fcc1eca983786b741ef61a69fd560a96ad Mon Sep 17 00:00:00 2001 From: Zishu Yu Date: Sun, 14 Jun 2026 03:08:05 +0800 Subject: [PATCH 33/40] build(headless): package Linux agentserver --- Makefile | 16 +++++- internal/vscode/linux_package_test.go | 69 ++++++++++++++++++++++++++ scripts/linux-package-common.sh | 71 +++++++++++++++++++++++++++ scripts/package-linux.sh | 68 +++++++++++++++++++++++++ 4 files changed, 223 insertions(+), 1 deletion(-) create mode 100644 internal/vscode/linux_package_test.go create mode 100644 scripts/linux-package-common.sh create mode 100755 scripts/package-linux.sh diff --git a/Makefile b/Makefile index 03982b5..4f1b5af 100644 --- a/Makefile +++ b/Makefile @@ -1,4 +1,4 @@ -.PHONY: all build test test-unit test-integration lint clean cross-windows package help ui-build ui-test +.PHONY: all build test test-unit test-integration lint clean cross-windows cross-linux package package-linux help ui-build ui-test GO ?= go GOFLAGS ?= -trimpath @@ -14,6 +14,7 @@ all: build help: @echo "make build - build native binaries to dist//" @echo "make cross-windows - cross-compile windows/amd64 to dist/windows/ (depends on ui-build)" + @echo "make cross-linux - cross-compile linux amd64/arm64 agentserver" @echo "make test - go test -race ./..." @echo "make test-unit - unit tests only (-short)" @echo "make test-integration - integration tests (test/integration)" @@ -22,6 +23,7 @@ help: @echo "make ui-build - build onboarding Vue front-end into internal/ui/assets/dist/" @echo "make ui-test - run frontend unit tests" @echo "make package - build Windows .exe installer (requires Inno Setup; depends on ui-build + ext-build)" + @echo "make package-linux - build Linux headless tarballs" @echo "make clean - rm dist/ and out/" build: ui-build @@ -43,6 +45,15 @@ cross-windows: ui-build -o $(DIST)/windows/$$cmd.exe ./cmd/$$cmd ; \ done +cross-linux: + @mkdir -p $(DIST)/linux/amd64 $(DIST)/linux/arm64 + @for arch in amd64 arm64; do \ + echo "==> cross-building agentserver (linux/$$arch)"; \ + GOOS=linux GOARCH=$$arch \ + $(GO) build $(GOFLAGS) -ldflags="$(LDFLAGS)" \ + -o $(DIST)/linux/$$arch/agentserver ./cmd/agentserver ; \ + done + test: ui-build $(GO) test -race -count=1 ./... @@ -69,5 +80,8 @@ ui-test: package: cross-windows ext-build bash scripts/package-windows.sh +package-linux: cross-linux + bash scripts/package-linux.sh + clean: rm -rf $(DIST) out coverage.out internal/ui/assets/dist diff --git a/internal/vscode/linux_package_test.go b/internal/vscode/linux_package_test.go new file mode 100644 index 0000000..5451c8d --- /dev/null +++ b/internal/vscode/linux_package_test.go @@ -0,0 +1,69 @@ +package vscode + +import ( + "os" + "strings" + "testing" +) + +func TestLinuxPackageCommonPinsLoomAssets(t *testing.T) { + body, err := os.ReadFile("../../scripts/linux-package-common.sh") + if err != nil { + t.Fatal(err) + } + text := string(body) + for _, want := range []string{ + `LOOM_RELEASE="v0.0.5"`, + `driver-agent.linux-amd64`, + `driver-agent.linux-arm64`, + `slave-agent.linux-amd64`, + `slave-agent.linux-arm64`, + `9dd94809801ff71d3e4c26581d48d44796c8e8be28be116b44d02cbd9fcb946c`, + `1c0a60bfb677a55159dea145dc46ead489b442d2cc55403dd451f3fadec4c7b5`, + `ce7d0b552a2ee880ef288d14c0d399630b961592fc73e78e98cece7a824ea965`, + `f7b0740cfb9d9a2c6fa1ad5f015b18c7ee4b3f618fe7082bb00bb828dc683ee6`, + `driver-skills.tar.gz`, + `driver-superpower-skills.tar.gz`, + `driver-codex-prompts.tar.gz`, + } { + if !strings.Contains(text, want) { + t.Fatalf("linux-package-common.sh missing %q", want) + } + } +} + +func TestPackageLinuxBuildsBothArchitectures(t *testing.T) { + body, err := os.ReadFile("../../scripts/package-linux.sh") + if err != nil { + t.Fatal(err) + } + text := string(body) + for _, want := range []string{ + `for arch in amd64 arm64`, + `GOOS=linux GOARCH="$arch"`, + `./cmd/agentserver`, + `agentserver-linux-$arch.tar.gz`, + `codex-manifest-linux-$arch.json`, + } { + if !strings.Contains(text, want) { + t.Fatalf("package-linux.sh missing %q", want) + } + } +} + +func TestMakefileExposesLinuxTargets(t *testing.T) { + body, err := os.ReadFile("../../Makefile") + if err != nil { + t.Fatal(err) + } + text := string(body) + for _, want := range []string{ + `cross-linux:`, + `package-linux: cross-linux`, + `bash scripts/package-linux.sh`, + } { + if !strings.Contains(text, want) { + t.Fatalf("Makefile missing %q", want) + } + } +} diff --git a/scripts/linux-package-common.sh b/scripts/linux-package-common.sh new file mode 100644 index 0000000..5912287 --- /dev/null +++ b/scripts/linux-package-common.sh @@ -0,0 +1,71 @@ +LOOM_RELEASE="v0.0.5" +LOOM_BASE_URL="https://github.com/agentserver/loom/releases/download/$LOOM_RELEASE" + +LOOM_DRIVER_AMD64_ASSET="driver-agent.linux-amd64" +LOOM_DRIVER_AMD64_SHA256="9dd94809801ff71d3e4c26581d48d44796c8e8be28be116b44d02cbd9fcb946c" +LOOM_DRIVER_ARM64_ASSET="driver-agent.linux-arm64" +LOOM_DRIVER_ARM64_SHA256="1c0a60bfb677a55159dea145dc46ead489b442d2cc55403dd451f3fadec4c7b5" + +LOOM_SLAVE_AMD64_ASSET="slave-agent.linux-amd64" +LOOM_SLAVE_AMD64_SHA256="ce7d0b552a2ee880ef288d14c0d399630b961592fc73e78e98cece7a824ea965" +LOOM_SLAVE_ARM64_ASSET="slave-agent.linux-arm64" +LOOM_SLAVE_ARM64_SHA256="f7b0740cfb9d9a2c6fa1ad5f015b18c7ee4b3f618fe7082bb00bb828dc683ee6" + +LOOM_DRIVER_SKILLS_ASSET="driver-skills.tar.gz" +LOOM_DRIVER_SKILLS_SHA256="4466b0342eaa90284dc4de0f0c03e6d08dbe02e4c12d0da6e7cb433c61ea1a0c" +SUPERPOWER_SKILLS_ASSET="driver-superpower-skills.tar.gz" +LOOM_DRIVER_CODEX_PROMPTS_ASSET="driver-codex-prompts.tar.gz" + +verify_sha256() { + local path expected sum + path="$1" + expected="$2" + [[ -f "$path" ]] || return 1 + sum=$(sha256sum "$path" | awk '{print $1}') + [[ "$sum" == "$expected" ]] +} + +download_loom_asset() { + local asset cache expected url sum + asset="$1" + cache="$2" + expected="$3" + url="$LOOM_BASE_URL/$asset" + if verify_sha256 "$cache" "$expected"; then + echo "$asset: $(stat -c%s "$cache") bytes (cached)" + return 0 + fi + mkdir -p "$(dirname "$cache")" + rm -f "$cache" "$cache.part" + echo "Fetching loom $asset ..." + echo " URL: $url" + if ! curl --fail --location --retry 2 --retry-delay 2 --output "$cache.part" "$url"; then + rm -f "$cache.part" + echo "ERROR: failed to download loom $asset" >&2 + exit 2 + fi + sum=$(sha256sum "$cache.part" | awk '{print $1}') + if [[ "$sum" != "$expected" ]]; then + rm -f "$cache.part" + echo "ERROR: loom $asset SHA256 mismatch: got $sum want $expected" >&2 + exit 2 + fi + mv "$cache.part" "$cache" + echo "$asset: $(stat -c%s "$cache") bytes (cached)" +} + +loom_asset_for_arch() { + local kind arch + kind="$1" + arch="$2" + case "$kind:$arch" in + driver:amd64) echo "$LOOM_DRIVER_AMD64_ASSET $LOOM_DRIVER_AMD64_SHA256" ;; + driver:arm64) echo "$LOOM_DRIVER_ARM64_ASSET $LOOM_DRIVER_ARM64_SHA256" ;; + slave:amd64) echo "$LOOM_SLAVE_AMD64_ASSET $LOOM_SLAVE_AMD64_SHA256" ;; + slave:arm64) echo "$LOOM_SLAVE_ARM64_ASSET $LOOM_SLAVE_ARM64_SHA256" ;; + *) + echo "ERROR: unsupported loom asset kind/arch: $kind/$arch" >&2 + return 2 + ;; + esac +} diff --git a/scripts/package-linux.sh b/scripts/package-linux.sh new file mode 100755 index 0000000..ff3f9d0 --- /dev/null +++ b/scripts/package-linux.sh @@ -0,0 +1,68 @@ +#!/usr/bin/env bash +set -euo pipefail + +ROOT="$(cd "$(dirname "${BASH_SOURCE[0]}")/.." && pwd)" +cd "$ROOT" + +VERSION="${VERSION:-0.1.2}" +OUT="${OUT:-dist}" +source scripts/linux-package-common.sh + +mkdir -p "$OUT/cache/loom/$LOOM_RELEASE" "$OUT/linux" + +loom_cache_path() { + local asset + asset="$1" + echo "$OUT/cache/loom/$LOOM_RELEASE/$asset" +} + +download_support_assets() { + download_loom_asset \ + "$LOOM_DRIVER_SKILLS_ASSET" \ + "$(loom_cache_path "$LOOM_DRIVER_SKILLS_ASSET")" \ + "$LOOM_DRIVER_SKILLS_SHA256" + python3 scripts/package-superpower-skills.py "$OUT/cache/superpowers/$SUPERPOWER_SKILLS_ASSET" + python3 scripts/package-driver-codex-prompts.py "$(loom_cache_path "$LOOM_DRIVER_CODEX_PROMPTS_ASSET")" +} + +download_arch_assets() { + local arch driver_asset driver_sha slave_asset slave_sha + arch="$1" + read -r driver_asset driver_sha < <(loom_asset_for_arch driver "$arch") + read -r slave_asset slave_sha < <(loom_asset_for_arch slave "$arch") + download_loom_asset "$driver_asset" "$(loom_cache_path "$driver_asset")" "$driver_sha" + download_loom_asset "$slave_asset" "$(loom_cache_path "$slave_asset")" "$slave_sha" +} + +stage_arch() { + local arch stage tarball driver_asset driver_sha slave_asset slave_sha + arch="$1" + stage="$OUT/linux/agentserver-linux-$arch" + tarball="$OUT/linux/agentserver-linux-$arch.tar.gz" + read -r driver_asset driver_sha < <(loom_asset_for_arch driver "$arch") + read -r slave_asset slave_sha < <(loom_asset_for_arch slave "$arch") + + rm -rf "$stage" + mkdir -p "$stage" + cp "$OUT/linux/$arch/agentserver" "$stage/agentserver" + cp "$(loom_cache_path "$driver_asset")" "$stage/driver-agent" + cp "$(loom_cache_path "$slave_asset")" "$stage/slave-agent" + cp "$(loom_cache_path "$LOOM_DRIVER_SKILLS_ASSET")" "$stage/driver-skills.tar.gz" + cp "$OUT/cache/superpowers/$SUPERPOWER_SKILLS_ASSET" "$stage/driver-superpower-skills.tar.gz" + cp "$(loom_cache_path "$LOOM_DRIVER_CODEX_PROMPTS_ASSET")" "$stage/driver-codex-prompts.tar.gz" + cp "packaging/linux/codex-manifest-linux-$arch.json" "$stage/codex-manifest-linux-$arch.json" + chmod 0755 "$stage/agentserver" "$stage/driver-agent" "$stage/slave-agent" + + rm -f "$tarball" "$tarball.sha256" + (cd "$OUT/linux" && tar -czf "agentserver-linux-$arch.tar.gz" "agentserver-linux-$arch") + sha256sum "$tarball" > "$tarball.sha256" + echo "wrote $tarball" +} + +download_support_assets +for arch in amd64 arm64; do + mkdir -p "$OUT/linux/$arch" + GOOS=linux GOARCH="$arch" go build -trimpath -ldflags="-s -w" -o "$OUT/linux/$arch/agentserver" ./cmd/agentserver + download_arch_assets "$arch" + stage_arch "$arch" +done From 507069f6f83b408c9d570979952c4074f6111845 Mon Sep 17 00:00:00 2001 From: Zishu Yu Date: Sun, 14 Jun 2026 03:15:44 +0800 Subject: [PATCH 34/40] fix(headless): package prebuilt Linux binaries --- internal/vscode/linux_package_test.go | 24 ++++++++++++++++++++++-- scripts/package-linux.sh | 8 +++++--- 2 files changed, 27 insertions(+), 5 deletions(-) diff --git a/internal/vscode/linux_package_test.go b/internal/vscode/linux_package_test.go index 5451c8d..912cc1c 100644 --- a/internal/vscode/linux_package_test.go +++ b/internal/vscode/linux_package_test.go @@ -40,8 +40,7 @@ func TestPackageLinuxBuildsBothArchitectures(t *testing.T) { text := string(body) for _, want := range []string{ `for arch in amd64 arm64`, - `GOOS=linux GOARCH="$arch"`, - `./cmd/agentserver`, + `run make cross-linux first`, `agentserver-linux-$arch.tar.gz`, `codex-manifest-linux-$arch.json`, } { @@ -51,6 +50,25 @@ func TestPackageLinuxBuildsBothArchitectures(t *testing.T) { } } +func TestPackageLinuxConsumesCrossLinuxBuildOutputs(t *testing.T) { + body, err := os.ReadFile("../../scripts/package-linux.sh") + if err != nil { + t.Fatal(err) + } + text := string(body) + if strings.Contains(text, "go build") { + t.Fatalf("package-linux.sh should package cross-linux outputs, not rebuild:\n%s", text) + } + for _, want := range []string{ + `cp "$OUT/linux/$arch/agentserver" "$stage/agentserver"`, + `sha256sum "agentserver-linux-$arch.tar.gz" > "agentserver-linux-$arch.tar.gz.sha256"`, + } { + if !strings.Contains(text, want) { + t.Fatalf("package-linux.sh missing %q", want) + } + } +} + func TestMakefileExposesLinuxTargets(t *testing.T) { body, err := os.ReadFile("../../Makefile") if err != nil { @@ -60,6 +78,8 @@ func TestMakefileExposesLinuxTargets(t *testing.T) { for _, want := range []string{ `cross-linux:`, `package-linux: cross-linux`, + `GOOS=linux GOARCH=$$arch`, + `./cmd/agentserver`, `bash scripts/package-linux.sh`, } { if !strings.Contains(text, want) { diff --git a/scripts/package-linux.sh b/scripts/package-linux.sh index ff3f9d0..95e7ca2 100755 --- a/scripts/package-linux.sh +++ b/scripts/package-linux.sh @@ -55,14 +55,16 @@ stage_arch() { rm -f "$tarball" "$tarball.sha256" (cd "$OUT/linux" && tar -czf "agentserver-linux-$arch.tar.gz" "agentserver-linux-$arch") - sha256sum "$tarball" > "$tarball.sha256" + (cd "$OUT/linux" && sha256sum "agentserver-linux-$arch.tar.gz" > "agentserver-linux-$arch.tar.gz.sha256") echo "wrote $tarball" } download_support_assets for arch in amd64 arm64; do - mkdir -p "$OUT/linux/$arch" - GOOS=linux GOARCH="$arch" go build -trimpath -ldflags="-s -w" -o "$OUT/linux/$arch/agentserver" ./cmd/agentserver + if [[ ! -x "$OUT/linux/$arch/agentserver" ]]; then + echo "ERROR: missing prebuilt $OUT/linux/$arch/agentserver; run make cross-linux first" >&2 + exit 2 + fi download_arch_assets "$arch" stage_arch "$arch" done From 059743290fe2c48078b948c8b45343ab7d6b21c8 Mon Sep 17 00:00:00 2001 From: Zishu Yu Date: Sun, 14 Jun 2026 03:22:59 +0800 Subject: [PATCH 35/40] fix(headless): align Linux packaging outputs --- Makefile | 2 +- internal/vscode/linux_package_test.go | 21 ++++++++++++++++++++- scripts/package-linux.sh | 15 +++++++++++---- 3 files changed, 32 insertions(+), 6 deletions(-) diff --git a/Makefile b/Makefile index 4f1b5af..472947d 100644 --- a/Makefile +++ b/Makefile @@ -81,7 +81,7 @@ package: cross-windows ext-build bash scripts/package-windows.sh package-linux: cross-linux - bash scripts/package-linux.sh + OUT="$(DIST)" bash scripts/package-linux.sh clean: rm -rf $(DIST) out coverage.out internal/ui/assets/dist diff --git a/internal/vscode/linux_package_test.go b/internal/vscode/linux_package_test.go index 912cc1c..3ab1ae5 100644 --- a/internal/vscode/linux_package_test.go +++ b/internal/vscode/linux_package_test.go @@ -69,6 +69,25 @@ func TestPackageLinuxConsumesCrossLinuxBuildOutputs(t *testing.T) { } } +func TestPackageLinuxPreflightsBinariesBeforeDownloads(t *testing.T) { + body, err := os.ReadFile("../../scripts/package-linux.sh") + if err != nil { + t.Fatal(err) + } + text := string(body) + preflight := strings.Index(text, "preflight_binaries") + downloads := strings.Index(text, "download_support_assets") + if preflight < 0 { + t.Fatalf("package-linux.sh missing preflight_binaries") + } + if downloads < 0 { + t.Fatalf("package-linux.sh missing download_support_assets") + } + if preflight > downloads { + t.Fatalf("package-linux.sh should preflight binaries before downloads") + } +} + func TestMakefileExposesLinuxTargets(t *testing.T) { body, err := os.ReadFile("../../Makefile") if err != nil { @@ -80,7 +99,7 @@ func TestMakefileExposesLinuxTargets(t *testing.T) { `package-linux: cross-linux`, `GOOS=linux GOARCH=$$arch`, `./cmd/agentserver`, - `bash scripts/package-linux.sh`, + `OUT="$(DIST)" bash scripts/package-linux.sh`, } { if !strings.Contains(text, want) { t.Fatalf("Makefile missing %q", want) diff --git a/scripts/package-linux.sh b/scripts/package-linux.sh index 95e7ca2..36a2e64 100755 --- a/scripts/package-linux.sh +++ b/scripts/package-linux.sh @@ -16,6 +16,16 @@ loom_cache_path() { echo "$OUT/cache/loom/$LOOM_RELEASE/$asset" } +preflight_binaries() { + local arch + for arch in amd64 arm64; do + if [[ ! -x "$OUT/linux/$arch/agentserver" ]]; then + echo "ERROR: missing prebuilt $OUT/linux/$arch/agentserver; run make cross-linux first" >&2 + exit 2 + fi + done +} + download_support_assets() { download_loom_asset \ "$LOOM_DRIVER_SKILLS_ASSET" \ @@ -59,12 +69,9 @@ stage_arch() { echo "wrote $tarball" } +preflight_binaries download_support_assets for arch in amd64 arm64; do - if [[ ! -x "$OUT/linux/$arch/agentserver" ]]; then - echo "ERROR: missing prebuilt $OUT/linux/$arch/agentserver; run make cross-linux first" >&2 - exit 2 - fi download_arch_assets "$arch" stage_arch "$arch" done From 88f10211356522c38512f68637df0bed6a8fd2b8 Mon Sep 17 00:00:00 2001 From: Zishu Yu Date: Sun, 14 Jun 2026 03:26:32 +0800 Subject: [PATCH 36/40] test(headless): assert Linux packaging preflight order --- internal/vscode/linux_package_test.go | 5 ++++- 1 file changed, 4 insertions(+), 1 deletion(-) diff --git a/internal/vscode/linux_package_test.go b/internal/vscode/linux_package_test.go index 3ab1ae5..9c53733 100644 --- a/internal/vscode/linux_package_test.go +++ b/internal/vscode/linux_package_test.go @@ -84,7 +84,10 @@ func TestPackageLinuxPreflightsBinariesBeforeDownloads(t *testing.T) { t.Fatalf("package-linux.sh missing download_support_assets") } if preflight > downloads { - t.Fatalf("package-linux.sh should preflight binaries before downloads") + t.Fatalf("package-linux.sh should define preflight before downloads") + } + if !strings.Contains(text, "preflight_binaries\ndownload_support_assets") { + t.Fatalf("package-linux.sh should call preflight_binaries immediately before download_support_assets") } } From 964bb867e1ba193f516c7a75b4bc2c69e6a2095e Mon Sep 17 00:00:00 2001 From: Zishu Yu Date: Sun, 14 Jun 2026 03:45:42 +0800 Subject: [PATCH 37/40] fix(headless): harden MCP access and token refresh --- cmd/agentserver/main.go | 23 +++++--- cmd/agentserver/main_test.go | 19 ++++++ internal/headless/driver.go | 3 + internal/headless/driver_test.go | 40 +++++++++++++ internal/modelaccess/access.go | 23 +++++++- internal/modelaccess/access_test.go | 90 +++++++++++++++++++++++++++++ 6 files changed, 190 insertions(+), 8 deletions(-) diff --git a/cmd/agentserver/main.go b/cmd/agentserver/main.go index 286697e..6e1ebba 100644 --- a/cmd/agentserver/main.go +++ b/cmd/agentserver/main.go @@ -22,6 +22,7 @@ import ( type app struct { ensureAccess func(context.Context) error + ensureMCPAccess func(context.Context) error ensureCodex func(context.Context) error runSlave func(context.Context) error installDriver func(context.Context) error @@ -54,14 +55,13 @@ func newApp() app { pkg := headless.PackagePaths(exe) sec := secrets.New(p.SecretsFile) cachedCodex := "" - - return app{ - ensureAccess: func(ctx context.Context) error { + ensureModelAccess := func(out io.Writer) func(context.Context) error { + return func(ctx context.Context) error { if _, err := modelaccess.Ensure(ctx, modelaccess.EnsureOptions{ CodexConfigPath: p.CodexConfigFile, Secrets: sec, PrintChallenge: func(title string, ch oauth.DeviceCodeChallenge) { - terminalauth.PrintChallenge(os.Stdout, title, ch, terminalauth.DefaultQR) + terminalauth.PrintChallenge(out, title, ch, terminalauth.DefaultQR) }, StartDaemon: func(ctx context.Context) error { return modelaccess.EnsureDaemon(ctx, modelaccess.EnsureDaemonOptions{ @@ -73,7 +73,12 @@ func newApp() app { return err } return nil - }, + } + } + + return app{ + ensureAccess: ensureModelAccess(os.Stdout), + ensureMCPAccess: ensureModelAccess(os.Stderr), ensureCodex: func(ctx context.Context) error { codexRuntime, err := headless.ResolveCodex(ctx, headless.CodexResolveOptions{ Paths: p, @@ -148,8 +153,12 @@ func (a app) run(ctx context.Context, args []string) error { return fmt.Errorf("unknown command %q", cmd) } - if a.ensureAccess != nil { - if err := a.ensureAccess(ctx); err != nil { + ensureAccess := a.ensureAccess + if cmd == "serve-driver-mcp" && a.ensureMCPAccess != nil { + ensureAccess = a.ensureMCPAccess + } + if ensureAccess != nil { + if err := ensureAccess(ctx); err != nil { return err } } diff --git a/cmd/agentserver/main_test.go b/cmd/agentserver/main_test.go index 6c586c1..1578d6f 100644 --- a/cmd/agentserver/main_test.go +++ b/cmd/agentserver/main_test.go @@ -123,6 +123,25 @@ func TestRunCommandResolvesCodexOnlyForDefaultSlave(t *testing.T) { } } +func TestRunCommandUsesMCPAccessForServeDriverMCP(t *testing.T) { + access := 0 + mcpAccess := 0 + app := app{ + ensureAccess: func(context.Context) error { access++; return nil }, + ensureMCPAccess: func(context.Context) error { mcpAccess++; return nil }, + serveDriverMCP: func(context.Context) error { return nil }, + } + if err := app.run(context.Background(), []string{"serve-driver-mcp"}); err != nil { + t.Fatalf("run: %v", err) + } + if access != 0 { + t.Fatalf("ensureAccess calls=%d, want 0", access) + } + if mcpAccess != 1 { + t.Fatalf("ensureMCPAccess calls=%d, want 1", mcpAccess) + } +} + func TestRunCommandSkipsAccessForDaemon(t *testing.T) { access := 0 daemonCalls := 0 diff --git a/internal/headless/driver.go b/internal/headless/driver.go index 9f0af74..08403de 100644 --- a/internal/headless/driver.go +++ b/internal/headless/driver.go @@ -84,6 +84,9 @@ func ServeDriverMCP(ctx context.Context, opts DriverMCPOptions) error { if err := writeDriverConfig(sessionConfig, st, proxyToken, tunnelToken, workDir); err != nil { return err } + defer func() { + _ = os.Remove(sessionConfig) + }() run := opts.Exec if run == nil { run = execDriverMCP diff --git a/internal/headless/driver_test.go b/internal/headless/driver_test.go index ef33403..d4019d1 100644 --- a/internal/headless/driver_test.go +++ b/internal/headless/driver_test.go @@ -159,6 +159,46 @@ func TestServeDriverMCPWritesSessionConfigWithCurrentWorkdir(t *testing.T) { } } +func TestServeDriverMCPRemovesSessionConfigAfterExit(t *testing.T) { + temp := t.TempDir() + p := testDriverPaths(temp) + pkg := testDriverPackage(temp) + sec := secrets.New(p.SecretsFile) + if err := sec.Set("agentserver_ws_api_key", "proxy-token"); err != nil { + t.Fatal(err) + } + if err := sec.Set("agentserver_tunnel_token", "tunnel-token"); err != nil { + t.Fatal(err) + } + writeDriverState(t, p, state.AgentserverState{ + BaseURL: "https://agent.test", + SandboxID: "sb-1", + ShortID: "abc123", + WorkspaceID: "ws-1", + }) + var sessionConfig string + + err := ServeDriverMCP(context.Background(), DriverMCPOptions{ + Paths: p, + Package: pkg, + Secrets: sec, + WorkDir: temp, + Exec: func(_ context.Context, _ string, args []string) error { + sessionConfig = args[2] + if _, err := os.Stat(sessionConfig); err != nil { + t.Fatalf("session config missing during Exec: %v", err) + } + return nil + }, + }) + if err != nil { + t.Fatalf("ServeDriverMCP: %v", err) + } + if _, err := os.Stat(sessionConfig); !errors.Is(err, os.ErrNotExist) { + t.Fatalf("session config stat after exit = %v, want not exist", err) + } +} + func TestSwitchWorkspaceForcesDeviceLogin(t *testing.T) { t.Run("empty state", func(t *testing.T) { temp := t.TempDir() diff --git a/internal/modelaccess/access.go b/internal/modelaccess/access.go index a9b500f..3ad89cc 100644 --- a/internal/modelaccess/access.go +++ b/internal/modelaccess/access.go @@ -15,6 +15,8 @@ import ( "github.com/agentserver/agentserver-pkg/internal/tokenrefresh" ) +const proxyRefreshBefore = 30 * time.Minute + type Mode string const ( @@ -119,7 +121,26 @@ func ensureProxyCredentials(ctx context.Context, opts EnsureOptions) error { } return err } - return nil + if accessTokenFresh(opts.Secrets, opts.Now()) { + return nil + } + err = refreshOnce(ctx, opts) + if tokenrefresh.ReauthRequired(err) { + return runDeviceLogin(ctx, opts) + } + return err +} + +func accessTokenFresh(sec secrets.Store, now time.Time) bool { + raw, err := sec.Get(tokenrefresh.AccessTokenExpiresAtKey) + if err != nil || raw == "" { + return false + } + expiresAt, err := time.Parse(time.RFC3339, raw) + if err != nil { + return false + } + return expiresAt.After(now.Add(proxyRefreshBefore)) } func refreshOnce(ctx context.Context, opts EnsureOptions) error { diff --git a/internal/modelaccess/access_test.go b/internal/modelaccess/access_test.go index e4685ad..d0a60c8 100644 --- a/internal/modelaccess/access_test.go +++ b/internal/modelaccess/access_test.go @@ -145,6 +145,7 @@ func TestEnsureProxyModeUsesExistingRefreshTokenWithoutPrompt(t *testing.T) { sec := secrets.New(filepath.Join(tmp, "secrets.json")) mustSetSecret(t, sec, tokenrefresh.AccessTokenKey, "existing-access") mustSetSecret(t, sec, tokenrefresh.RefreshTokenKey, "existing-refresh") + mustSetSecret(t, sec, tokenrefresh.AccessTokenExpiresAtKey, fixedNow().Add(time.Hour).Format(time.RFC3339)) var daemonStarted bool result, err := Ensure(context.Background(), EnsureOptions{ @@ -163,6 +164,7 @@ func TestEnsureProxyModeUsesExistingRefreshTokenWithoutPrompt(t *testing.T) { daemonStarted = true return nil }, + Now: fixedNow, }) if err != nil { t.Fatalf("Ensure() error = %v", err) @@ -177,6 +179,86 @@ func TestEnsureProxyModeUsesExistingRefreshTokenWithoutPrompt(t *testing.T) { assertSecret(t, sec, tokenrefresh.RefreshTokenKey, "existing-refresh") } +func TestEnsureProxyModeRefreshesExpiredAccessTokenBeforeReturning(t *testing.T) { + tmp := t.TempDir() + sec := secrets.New(filepath.Join(tmp, "secrets.json")) + mustSetSecret(t, sec, tokenrefresh.AccessTokenKey, "expired-access") + mustSetSecret(t, sec, tokenrefresh.RefreshTokenKey, "existing-refresh") + mustSetSecret(t, sec, tokenrefresh.AccessTokenExpiresAtKey, fixedNow().Add(-time.Minute).Format(time.RFC3339)) + var refreshed, daemonStarted bool + + result, err := Ensure(context.Background(), EnsureOptions{ + CodexConfigPath: filepath.Join(tmp, "codex", "config.toml"), + Secrets: sec, + Env: func(string) string { return "" }, + Refresh: func(_ context.Context, _ oauth.AuthCodeConfig, refreshToken string) (oauth.Token, error) { + refreshed = true + if refreshToken != "existing-refresh" { + t.Fatalf("refresh token = %q, want existing-refresh", refreshToken) + } + return oauth.Token{AccessToken: "fresh-access", ExpiresIn: 3600}, nil + }, + RequestDeviceCode: func(context.Context, oauth.Config) (oauth.DeviceCodeChallenge, error) { + t.Fatal("RequestDeviceCode called") + return oauth.DeviceCodeChallenge{}, nil + }, + StartDaemon: func(context.Context) error { + daemonStarted = true + return nil + }, + Now: fixedNow, + }) + if err != nil { + t.Fatalf("Ensure() error = %v", err) + } + if result.Mode != ModeLocalProxy { + t.Fatalf("Mode = %q, want %q", result.Mode, ModeLocalProxy) + } + if !refreshed || !daemonStarted { + t.Fatalf("calls: refresh=%v daemon=%v", refreshed, daemonStarted) + } + assertSecret(t, sec, tokenrefresh.AccessTokenKey, "fresh-access") +} + +func TestEnsureProxyModeRunsDeviceLoginWhenExpiredAccessRefreshNeedsReauth(t *testing.T) { + tmp := t.TempDir() + sec := secrets.New(filepath.Join(tmp, "secrets.json")) + mustSetSecret(t, sec, tokenrefresh.AccessTokenKey, "expired-access") + mustSetSecret(t, sec, tokenrefresh.RefreshTokenKey, "existing-refresh") + mustSetSecret(t, sec, tokenrefresh.AccessTokenExpiresAtKey, fixedNow().Add(-time.Minute).Format(time.RFC3339)) + var refreshed, requestedDeviceCode, daemonStarted bool + + _, err := Ensure(context.Background(), EnsureOptions{ + CodexConfigPath: filepath.Join(tmp, "codex", "config.toml"), + Secrets: sec, + Env: func(string) string { return "" }, + Refresh: func(_ context.Context, _ oauth.AuthCodeConfig, refreshToken string) (oauth.Token, error) { + refreshed = true + return oauth.Token{}, fmt.Errorf("refresh failed: %w", oauth.ErrInvalidGrant) + }, + RequestDeviceCode: func(context.Context, oauth.Config) (oauth.DeviceCodeChallenge, error) { + requestedDeviceCode = true + return oauth.DeviceCodeChallenge{DeviceCode: "device", UserCode: "code", ExpiresIn: 600}, nil + }, + PrintChallenge: func(string, oauth.DeviceCodeChallenge) {}, + PollToken: func(context.Context, oauth.Config, oauth.DeviceCodeChallenge) (oauth.Token, error) { + return oauth.Token{AccessToken: "device-access", RefreshToken: "device-refresh", ExpiresIn: 3600}, nil + }, + StartDaemon: func(context.Context) error { + daemonStarted = true + return nil + }, + Now: fixedNow, + }) + if err != nil { + t.Fatalf("Ensure() error = %v", err) + } + if !refreshed || !requestedDeviceCode || !daemonStarted { + t.Fatalf("calls: refresh=%v request=%v daemon=%v", refreshed, requestedDeviceCode, daemonStarted) + } + assertSecret(t, sec, tokenrefresh.AccessTokenKey, "device-access") +} + func TestEnsureProxyModeRefreshesWhenAccessTokenMissing(t *testing.T) { tmp := t.TempDir() sec := secrets.New(filepath.Join(tmp, "secrets.json")) @@ -367,6 +449,7 @@ func TestEnsureProxyModeDoesNotWriteProxyConfigWhenSetEnvFails(t *testing.T) { sec := secrets.New(filepath.Join(tmp, "secrets.json")) mustSetSecret(t, sec, tokenrefresh.AccessTokenKey, "existing-access") mustSetSecret(t, sec, tokenrefresh.RefreshTokenKey, "existing-refresh") + mustSetSecret(t, sec, tokenrefresh.AccessTokenExpiresAtKey, fixedNow().Add(time.Hour).Format(time.RFC3339)) setEnvErr := errors.New("set env failed") _, err := Ensure(context.Background(), EnsureOptions{ @@ -384,6 +467,7 @@ func TestEnsureProxyModeDoesNotWriteProxyConfigWhenSetEnvFails(t *testing.T) { t.Fatal("StartDaemon called after SetEnv failure") return nil }, + Now: fixedNow, }) if !errors.Is(err, setEnvErr) { t.Fatalf("Ensure() error = %v, want %v", err, setEnvErr) @@ -398,6 +482,7 @@ func TestEnsureProxyModeDoesNotWriteProxyConfigWhenPersistEnvFails(t *testing.T) sec := secrets.New(filepath.Join(tmp, "secrets.json")) mustSetSecret(t, sec, tokenrefresh.AccessTokenKey, "existing-access") mustSetSecret(t, sec, tokenrefresh.RefreshTokenKey, "existing-refresh") + mustSetSecret(t, sec, tokenrefresh.AccessTokenExpiresAtKey, fixedNow().Add(time.Hour).Format(time.RFC3339)) persistErr := errors.New("persist env failed") _, err := Ensure(context.Background(), EnsureOptions{ @@ -412,6 +497,7 @@ func TestEnsureProxyModeDoesNotWriteProxyConfigWhenPersistEnvFails(t *testing.T) t.Fatal("StartDaemon called after PersistEnv failure") return nil }, + Now: fixedNow, }) if !errors.Is(err, persistErr) { t.Fatalf("Ensure() error = %v, want %v", err, persistErr) @@ -426,6 +512,7 @@ func TestEnsureProxyModeDoesNotWriteProxyConfigWhenStartDaemonFails(t *testing.T sec := secrets.New(filepath.Join(tmp, "secrets.json")) mustSetSecret(t, sec, tokenrefresh.AccessTokenKey, "existing-access") mustSetSecret(t, sec, tokenrefresh.RefreshTokenKey, "existing-refresh") + mustSetSecret(t, sec, tokenrefresh.AccessTokenExpiresAtKey, fixedNow().Add(time.Hour).Format(time.RFC3339)) daemonErr := errors.New("daemon failed") _, err := Ensure(context.Background(), EnsureOptions{ @@ -437,6 +524,7 @@ func TestEnsureProxyModeDoesNotWriteProxyConfigWhenStartDaemonFails(t *testing.T StartDaemon: func(context.Context) error { return daemonErr }, + Now: fixedNow, }) if !errors.Is(err, daemonErr) { t.Fatalf("Ensure() error = %v, want %v", err, daemonErr) @@ -449,6 +537,7 @@ func TestProxySettingsWrittenInProxyMode(t *testing.T) { sec := secrets.New(filepath.Join(tmp, "secrets.json")) mustSetSecret(t, sec, tokenrefresh.AccessTokenKey, "existing-access") mustSetSecret(t, sec, tokenrefresh.RefreshTokenKey, "existing-refresh") + mustSetSecret(t, sec, tokenrefresh.AccessTokenExpiresAtKey, fixedNow().Add(time.Hour).Format(time.RFC3339)) var setEnvKey, setEnvValue, persistEnvKey, persistEnvValue string var daemonStarted bool @@ -468,6 +557,7 @@ func TestProxySettingsWrittenInProxyMode(t *testing.T) { daemonStarted = true return nil }, + Now: fixedNow, }) if err != nil { t.Fatalf("Ensure() error = %v", err) From a39febdb4335ca8983ac03604094a12a75cd37e5 Mon Sep 17 00:00:00 2001 From: Zishu Yu Date: Sun, 14 Jun 2026 04:10:00 +0800 Subject: [PATCH 38/40] fix(headless): satisfy Linux auth and daemon review --- cmd/agentctl/cmd_test_subcommands_test.go | 4 +- cmd/launcher/main_test.go | 6 +- cmd/open-folder/main_test.go | 4 +- internal/codex/config.go | 74 +++++++++--- internal/codex/config_test.go | 45 +++++++- internal/headless/slave.go | 4 +- internal/headless/slave_test.go | 40 ++++++- internal/modelaccess/access.go | 16 +-- internal/modelaccess/access_test.go | 106 ++++++++++++++---- internal/modelaccess/daemon.go | 10 +- internal/modelaccess/daemon_process_unix.go | 21 ++++ .../modelaccess/daemon_process_windows.go | 17 +++ internal/modelaccess/daemon_unix_test.go | 45 ++++++++ internal/modelserver/device.go | 4 +- internal/modelserver/device_test.go | 11 +- internal/ui/orchestrator_real_test.go | 4 +- 16 files changed, 343 insertions(+), 68 deletions(-) create mode 100644 internal/modelaccess/daemon_process_unix.go create mode 100644 internal/modelaccess/daemon_process_windows.go create mode 100644 internal/modelaccess/daemon_unix_test.go diff --git a/cmd/agentctl/cmd_test_subcommands_test.go b/cmd/agentctl/cmd_test_subcommands_test.go index 71d496d..e290792 100644 --- a/cmd/agentctl/cmd_test_subcommands_test.go +++ b/cmd/agentctl/cmd_test_subcommands_test.go @@ -53,8 +53,8 @@ func TestOpenTestFolderCodexDesktopUsesDeepLinkAndWritesConfig(t *testing.T) { if !strings.Contains(string(b), `base_url = "`+modelproxy.DefaultBaseURL+`"`) { t.Fatalf("config missing local proxy base_url:\n%s", b) } - if !strings.Contains(string(b), `env_key = "`+codex.LocalProxyAPIKeyEnv+`"`) { - t.Fatalf("config missing local proxy env_key:\n%s", b) + if !strings.Contains(string(b), `experimental_bearer_token = "`+codex.LocalProxyAPIKeyValue+`"`) { + t.Fatalf("config missing local proxy bearer token:\n%s", b) } } diff --git a/cmd/launcher/main_test.go b/cmd/launcher/main_test.go index 397ca11..1d67d7e 100644 --- a/cmd/launcher/main_test.go +++ b/cmd/launcher/main_test.go @@ -893,7 +893,7 @@ func TestExecVSCodeEnsuresCodexConfigBeforeLaunch(t *testing.T) { for _, want := range []string{ `model_provider = "modelserver"`, `base_url = "` + modelproxy.DefaultBaseURL + `"`, - `env_key = "` + codex.LocalProxyAPIKeyEnv + `"`, + `experimental_bearer_token = "` + codex.LocalProxyAPIKeyValue + `"`, `[windows]`, `sandbox = "unelevated"`, } { @@ -992,8 +992,8 @@ func TestLaunchCompletedCodexDesktopWritesConfigAndOpensDeepLink(t *testing.T) { if !strings.Contains(string(b), `base_url = "`+modelproxy.DefaultBaseURL+`"`) { t.Fatalf("config missing local proxy base_url:\n%s", b) } - if !strings.Contains(string(b), `env_key = "`+codex.LocalProxyAPIKeyEnv+`"`) { - t.Fatalf("config missing local proxy env_key:\n%s", b) + if !strings.Contains(string(b), `experimental_bearer_token = "`+codex.LocalProxyAPIKeyValue+`"`) { + t.Fatalf("config missing local proxy bearer token:\n%s", b) } } diff --git a/cmd/open-folder/main_test.go b/cmd/open-folder/main_test.go index 1a30b0f..50ab700 100644 --- a/cmd/open-folder/main_test.go +++ b/cmd/open-folder/main_test.go @@ -63,7 +63,7 @@ func TestOpenFolderMigratesVSCodeSettingsBeforeLaunch(t *testing.T) { } for _, want := range []string{ `base_url = "` + modelproxy.DefaultBaseURL + `"`, - `env_key = "` + codex.LocalProxyAPIKeyEnv + `"`, + `experimental_bearer_token = "` + codex.LocalProxyAPIKeyValue + `"`, } { if !strings.Contains(string(b), want) { t.Fatalf("missing %q in:\n%s", want, b) @@ -94,7 +94,7 @@ func TestOpenFolderCodexDesktopUsesFolderDeepLink(t *testing.T) { } for _, want := range []string{ `base_url = "` + modelproxy.DefaultBaseURL + `"`, - `env_key = "` + codex.LocalProxyAPIKeyEnv + `"`, + `experimental_bearer_token = "` + codex.LocalProxyAPIKeyValue + `"`, } { if !strings.Contains(string(b), want) { t.Fatalf("missing %q in:\n%s", want, b) diff --git a/internal/codex/config.go b/internal/codex/config.go index 4ebbde8..728ea7a 100644 --- a/internal/codex/config.go +++ b/internal/codex/config.go @@ -16,16 +16,17 @@ import ( ) type Settings struct { - Provider string // e.g. "modelserver" - Model string // e.g. "gpt-5.5" - ModelReasoningEffort string // e.g. "high" - ApprovalsReviewer string // e.g. "guardian_subagent" - SandboxMode string // e.g. "danger-full-access" - WindowsSandbox string // e.g. "unelevated" - DeveloperInstructions string - BaseURL string // e.g. "https://code.ai.cs.ac.cn/v1" - EnvKey string // e.g. "OPENAI_API_KEY" - WireAPI string // e.g. "responses" + Provider string // e.g. "modelserver" + Model string // e.g. "gpt-5.5" + ModelReasoningEffort string // e.g. "high" + ApprovalsReviewer string // e.g. "guardian_subagent" + SandboxMode string // e.g. "danger-full-access" + WindowsSandbox string // e.g. "unelevated" + DeveloperInstructions string + BaseURL string // e.g. "https://code.ai.cs.ac.cn/v1" + EnvKey string // e.g. "OPENAI_API_KEY" + ExperimentalBearerToken string // e.g. "agentserver-local-proxy" + WireAPI string // e.g. "responses" } type MCPServer struct { @@ -56,7 +57,8 @@ const ( func ModelserverProxySettings(baseURL string) Settings { s := ModelserverSettings() s.BaseURL = baseURL - s.EnvKey = LocalProxyAPIKeyEnv + s.EnvKey = "" + s.ExperimentalBearerToken = LocalProxyAPIKeyValue return s } @@ -108,12 +110,18 @@ func UpdateConfig(path string, s Settings) error { if providers == nil { providers = map[string]any{} } - providers[s.Provider] = map[string]any{ + provider := map[string]any{ "name": s.Provider, "base_url": s.BaseURL, - "env_key": s.EnvKey, "wire_api": s.WireAPI, } + if s.EnvKey != "" { + provider["env_key"] = s.EnvKey + } + if s.ExperimentalBearerToken != "" { + provider["experimental_bearer_token"] = s.ExperimentalBearerToken + } + providers[s.Provider] = provider root["model_providers"] = providers var buf bytes.Buffer @@ -123,6 +131,46 @@ func UpdateConfig(path string, s Settings) error { return os.WriteFile(path, buf.Bytes(), 0o644) } +func HasModelserverDirectConfig(path string) (bool, error) { + b, err := os.ReadFile(path) + if err != nil { + if errors.Is(err, os.ErrNotExist) { + return false, nil + } + return false, fmt.Errorf("read config.toml: %w", err) + } + root := map[string]any{} + if _, err := toml.Decode(string(b), &root); err != nil { + return false, fmt.Errorf("parse existing config.toml: %w", err) + } + settings := ModelserverSettings() + if stringSetting(root, "model_provider") != settings.Provider { + return false, nil + } + providers := tableSetting(root, "model_providers") + if providers == nil { + return false, nil + } + provider := tableSetting(providers, settings.Provider) + if provider == nil { + return false, nil + } + return stringSetting(provider, "name") == settings.Provider && + stringSetting(provider, "base_url") == settings.BaseURL && + stringSetting(provider, "env_key") == settings.EnvKey && + stringSetting(provider, "wire_api") == settings.WireAPI, nil +} + +func tableSetting(root map[string]any, key string) map[string]any { + table, _ := root[key].(map[string]any) + return table +} + +func stringSetting(root map[string]any, key string) string { + value, _ := root[key].(string) + return strings.TrimSpace(value) +} + func UpdateMCPServer(path, name string, server MCPServer) error { if name == "" { return errors.New("MCP server name required") diff --git a/internal/codex/config_test.go b/internal/codex/config_test.go index ca165af..8280d84 100644 --- a/internal/codex/config_test.go +++ b/internal/codex/config_test.go @@ -52,12 +52,53 @@ func TestModelserverProxySettingsUsesStableLocalCredential(t *testing.T) { if got.BaseURL != "http://127.0.0.1:53452/v1" { t.Fatalf("BaseURL = %q, want local proxy URL", got.BaseURL) } - if got.EnvKey != LocalProxyAPIKeyEnv { - t.Fatalf("EnvKey = %q, want %q", got.EnvKey, LocalProxyAPIKeyEnv) + if got.EnvKey != "" { + t.Fatalf("EnvKey = %q, want empty local proxy env key", got.EnvKey) } if got.WireAPI != "responses" { t.Fatalf("WireAPI = %q, want responses", got.WireAPI) } + + dir := t.TempDir() + path := filepath.Join(dir, "config.toml") + if err := UpdateConfig(path, got); err != nil { + t.Fatal(err) + } + b, _ := os.ReadFile(path) + s := string(b) + if strings.Contains(s, "env_key") { + t.Fatalf("proxy config should not require an environment variable:\n%s", s) + } + if !strings.Contains(s, `experimental_bearer_token = "`+LocalProxyAPIKeyValue+`"`) { + t.Fatalf("proxy config missing stable bearer token:\n%s", s) + } +} + +func TestHasModelserverDirectConfigRequiresExactDirectProvider(t *testing.T) { + dir := t.TempDir() + path := filepath.Join(dir, "config.toml") + if ok, err := HasModelserverDirectConfig(path); err != nil || ok { + t.Fatalf("missing config: ok=%v err=%v, want false nil", ok, err) + } + writeCodexTestFile(t, path, strings.Join([]string{ + `model_provider = "modelserver"`, + ``, + `[model_providers.modelserver]`, + `name = "modelserver"`, + `base_url = "http://127.0.0.1:53452/v1"`, + `experimental_bearer_token = "agentserver-local-proxy"`, + `wire_api = "responses"`, + ``, + }, "\n")) + if ok, err := HasModelserverDirectConfig(path); err != nil || ok { + t.Fatalf("proxy config: ok=%v err=%v, want false nil", ok, err) + } + if err := UpdateConfig(path, ModelserverSettings()); err != nil { + t.Fatal(err) + } + if ok, err := HasModelserverDirectConfig(path); err != nil || !ok { + t.Fatalf("direct config: ok=%v err=%v, want true nil", ok, err) + } } func TestUpdateConfig_MergeKeepsOtherProvider(t *testing.T) { diff --git a/internal/headless/slave.go b/internal/headless/slave.go index 4403d82..ba58821 100644 --- a/internal/headless/slave.go +++ b/internal/headless/slave.go @@ -45,6 +45,8 @@ type SlaveProcessRunner interface { Run(context.Context, SlaveProcessRequest) error } +const defaultHeadlessSlaveName = "星池指挥官" + func RunSlave(ctx context.Context, opts SlaveOptions) error { stdout := opts.Stdout if stdout == nil { @@ -80,7 +82,7 @@ func RunSlave(ctx context.Context, opts SlaveOptions) error { } else if ok { name = existing.Name } else { - name = filepath.Base(canonicalFolder) + name = defaultHeadlessSlaveName if opts.NamePrompt != nil { prompted, err := opts.NamePrompt(name) if err != nil { diff --git a/internal/headless/slave_test.go b/internal/headless/slave_test.go index 2bcf8ae..7822716 100644 --- a/internal/headless/slave_test.go +++ b/internal/headless/slave_test.go @@ -34,8 +34,8 @@ func TestRunSlaveCreatesStableRegistryEntryAndWritesConfig(t *testing.T) { WorkDir: repo, ComputerName: "host", NamePrompt: func(defaultName string) (string, error) { - if defaultName != "repo" { - t.Fatalf("defaultName=%q, want repo", defaultName) + if defaultName != "星池指挥官" { + t.Fatalf("defaultName=%q, want 星池指挥官", defaultName) } return "worker", nil }, @@ -78,6 +78,42 @@ func TestRunSlaveCreatesStableRegistryEntryAndWritesConfig(t *testing.T) { } } +func TestRunSlaveDefaultsNewEntryToCommanderName(t *testing.T) { + temp := t.TempDir() + repo := filepath.Join(temp, "repo") + if err := os.Mkdir(repo, 0o755); err != nil { + t.Fatal(err) + } + runner := &fakeSlaveProcessRunner{} + + err := RunSlave(context.Background(), SlaveOptions{ + Paths: testSlavePaths(temp), + Package: Package{SlaveAgent: filepath.Join(temp, "pkg", "slave-agent")}, + WorkDir: repo, + ComputerName: "host", + NamePrompt: func(defaultName string) (string, error) { + if defaultName != "星池指挥官" { + t.Fatalf("defaultName=%q, want 星池指挥官", defaultName) + } + return "", nil + }, + Runner: runner, + Stdout: ioDiscard{}, + }) + if err != nil { + t.Fatalf("RunSlave: %v", err) + } + + configBody, err := os.ReadFile(runner.request.ConfigPath) + if err != nil { + t.Fatal(err) + } + config := string(configBody) + if !strings.Contains(config, "display_name: host-星池指挥官") { + t.Fatalf("config missing default commander display name:\n%s", config) + } +} + func TestRunSlaveReusesExistingEntryWithoutPrompt(t *testing.T) { temp := t.TempDir() repo := filepath.Join(temp, "repo") diff --git a/internal/modelaccess/access.go b/internal/modelaccess/access.go index 3ad89cc..b764639 100644 --- a/internal/modelaccess/access.go +++ b/internal/modelaccess/access.go @@ -54,10 +54,16 @@ func Ensure(ctx context.Context, opts EnsureOptions) (Result, error) { opts = defaultEnsureOptions(opts) if opts.Env(tokenrefresh.OpenAIAPIKeyEnv) != "" { - if err := codex.UpdateConfig(opts.CodexConfigPath, codex.ModelserverSettings()); err != nil { + directConfig, err := codex.HasModelserverDirectConfig(opts.CodexConfigPath) + if err != nil { return Result{}, err } - return Result{Mode: ModeDirectAPIKey}, nil + if directConfig { + if err := codex.UpdateConfig(opts.CodexConfigPath, codex.ModelserverSettings()); err != nil { + return Result{}, err + } + return Result{Mode: ModeDirectAPIKey}, nil + } } if opts.Secrets == nil { @@ -66,12 +72,6 @@ func Ensure(ctx context.Context, opts EnsureOptions) (Result, error) { if err := ensureProxyCredentials(ctx, opts); err != nil { return Result{}, err } - if err := opts.SetEnv(codex.LocalProxyAPIKeyEnv, codex.LocalProxyAPIKeyValue); err != nil { - return Result{}, err - } - if err := opts.PersistEnv(codex.LocalProxyAPIKeyEnv, codex.LocalProxyAPIKeyValue); err != nil { - return Result{}, err - } if opts.StartDaemon != nil { if err := opts.StartDaemon(ctx); err != nil { return Result{}, err diff --git a/internal/modelaccess/access_test.go b/internal/modelaccess/access_test.go index d0a60c8..5d9a0a4 100644 --- a/internal/modelaccess/access_test.go +++ b/internal/modelaccess/access_test.go @@ -19,11 +19,13 @@ import ( func TestEnsurePrefersLongLivedAPIKey(t *testing.T) { tmp := t.TempDir() + configPath := filepath.Join(tmp, "codex", "config.toml") + writeDirectConfig(t, configPath) sec := secrets.New(filepath.Join(tmp, "secrets.json")) var daemonStarted bool result, err := Ensure(context.Background(), EnsureOptions{ - CodexConfigPath: filepath.Join(tmp, "codex", "config.toml"), + CodexConfigPath: configPath, Secrets: sec, Env: func(key string) string { if key == tokenrefresh.OpenAIAPIKeyEnv { @@ -53,7 +55,7 @@ func TestEnsurePrefersLongLivedAPIKey(t *testing.T) { if daemonStarted { t.Fatal("StartDaemon called in direct API key mode") } - assertConfigContains(t, filepath.Join(tmp, "codex", "config.toml"), + assertConfigContains(t, configPath, `model_provider = "modelserver"`, `base_url = "https://code.ai.cs.ac.cn/v1"`, `env_key = "OPENAI_API_KEY"`, @@ -63,6 +65,7 @@ func TestEnsurePrefersLongLivedAPIKey(t *testing.T) { func TestEnsureDirectAPIKeyModeDoesNotRequireSecrets(t *testing.T) { tmp := t.TempDir() configPath := filepath.Join(tmp, "codex", "config.toml") + writeDirectConfig(t, configPath) result, err := Ensure(context.Background(), EnsureOptions{ CodexConfigPath: configPath, @@ -86,6 +89,50 @@ func TestEnsureDirectAPIKeyModeDoesNotRequireSecrets(t *testing.T) { ) } +func TestEnsureOpenAIAPIKeyWithoutMatchingDirectConfigUsesProxy(t *testing.T) { + tmp := t.TempDir() + configPath := filepath.Join(tmp, "codex", "config.toml") + sec := secrets.New(filepath.Join(tmp, "secrets.json")) + mustSetSecret(t, sec, tokenrefresh.AccessTokenKey, "existing-access") + mustSetSecret(t, sec, tokenrefresh.RefreshTokenKey, "existing-refresh") + mustSetSecret(t, sec, tokenrefresh.AccessTokenExpiresAtKey, fixedNow().Add(time.Hour).Format(time.RFC3339)) + var daemonStarted bool + + result, err := Ensure(context.Background(), EnsureOptions{ + CodexConfigPath: configPath, + Secrets: sec, + Env: func(key string) string { + if key == tokenrefresh.OpenAIAPIKeyEnv { + return "sk-test" + } + return "" + }, + RequestDeviceCode: func(context.Context, oauth.Config) (oauth.DeviceCodeChallenge, error) { + t.Fatal("RequestDeviceCode called despite fresh proxy credentials") + return oauth.DeviceCodeChallenge{}, nil + }, + StartDaemon: func(context.Context) error { + daemonStarted = true + return nil + }, + Now: fixedNow, + }) + if err != nil { + t.Fatalf("Ensure() error = %v", err) + } + if result.Mode != ModeLocalProxy { + t.Fatalf("Mode = %q, want %q", result.Mode, ModeLocalProxy) + } + if !daemonStarted { + t.Fatal("StartDaemon was not called") + } + assertConfigContains(t, configPath, + `base_url = "`+modelproxy.DefaultBaseURL+`"`, + `experimental_bearer_token = "`+codex.LocalProxyAPIKeyValue+`"`, + ) + assertConfigNotContains(t, configPath, `env_key = "OPENAI_API_KEY"`) +} + func TestEnsureProxyModeRunsDeviceLoginWhenRefreshMissing(t *testing.T) { tmp := t.TempDir() sec := secrets.New(filepath.Join(tmp, "secrets.json")) @@ -442,7 +489,7 @@ func TestEnsureRequiresSecrets(t *testing.T) { } } -func TestEnsureProxyModeDoesNotWriteProxyConfigWhenSetEnvFails(t *testing.T) { +func TestEnsureProxyModeIgnoresSetEnvFailure(t *testing.T) { tmp := t.TempDir() configPath := filepath.Join(tmp, "codex", "config.toml") writeDirectConfig(t, configPath) @@ -451,8 +498,9 @@ func TestEnsureProxyModeDoesNotWriteProxyConfigWhenSetEnvFails(t *testing.T) { mustSetSecret(t, sec, tokenrefresh.RefreshTokenKey, "existing-refresh") mustSetSecret(t, sec, tokenrefresh.AccessTokenExpiresAtKey, fixedNow().Add(time.Hour).Format(time.RFC3339)) setEnvErr := errors.New("set env failed") + var daemonStarted bool - _, err := Ensure(context.Background(), EnsureOptions{ + result, err := Ensure(context.Background(), EnsureOptions{ CodexConfigPath: configPath, Secrets: sec, Env: func(string) string { return "" }, @@ -464,18 +512,24 @@ func TestEnsureProxyModeDoesNotWriteProxyConfigWhenSetEnvFails(t *testing.T) { return nil }, StartDaemon: func(context.Context) error { - t.Fatal("StartDaemon called after SetEnv failure") + daemonStarted = true return nil }, Now: fixedNow, }) - if !errors.Is(err, setEnvErr) { - t.Fatalf("Ensure() error = %v, want %v", err, setEnvErr) + if err != nil { + t.Fatalf("Ensure() error = %v", err) } - assertConfigNotContains(t, configPath, `base_url = "`+modelproxy.DefaultBaseURL+`"`) + if result.Mode != ModeLocalProxy { + t.Fatalf("Mode = %q, want %q", result.Mode, ModeLocalProxy) + } + if !daemonStarted { + t.Fatal("StartDaemon was not called") + } + assertConfigContains(t, configPath, `base_url = "`+modelproxy.DefaultBaseURL+`"`) } -func TestEnsureProxyModeDoesNotWriteProxyConfigWhenPersistEnvFails(t *testing.T) { +func TestEnsureProxyModeIgnoresPersistEnvFailure(t *testing.T) { tmp := t.TempDir() configPath := filepath.Join(tmp, "codex", "config.toml") writeDirectConfig(t, configPath) @@ -484,8 +538,9 @@ func TestEnsureProxyModeDoesNotWriteProxyConfigWhenPersistEnvFails(t *testing.T) mustSetSecret(t, sec, tokenrefresh.RefreshTokenKey, "existing-refresh") mustSetSecret(t, sec, tokenrefresh.AccessTokenExpiresAtKey, fixedNow().Add(time.Hour).Format(time.RFC3339)) persistErr := errors.New("persist env failed") + var daemonStarted bool - _, err := Ensure(context.Background(), EnsureOptions{ + result, err := Ensure(context.Background(), EnsureOptions{ CodexConfigPath: configPath, Secrets: sec, Env: func(string) string { return "" }, @@ -494,15 +549,21 @@ func TestEnsureProxyModeDoesNotWriteProxyConfigWhenPersistEnvFails(t *testing.T) return persistErr }, StartDaemon: func(context.Context) error { - t.Fatal("StartDaemon called after PersistEnv failure") + daemonStarted = true return nil }, Now: fixedNow, }) - if !errors.Is(err, persistErr) { - t.Fatalf("Ensure() error = %v, want %v", err, persistErr) + if err != nil { + t.Fatalf("Ensure() error = %v", err) } - assertConfigNotContains(t, configPath, `base_url = "`+modelproxy.DefaultBaseURL+`"`) + if result.Mode != ModeLocalProxy { + t.Fatalf("Mode = %q, want %q", result.Mode, ModeLocalProxy) + } + if !daemonStarted { + t.Fatal("StartDaemon was not called") + } + assertConfigContains(t, configPath, `base_url = "`+modelproxy.DefaultBaseURL+`"`) } func TestEnsureProxyModeDoesNotWriteProxyConfigWhenStartDaemonFails(t *testing.T) { @@ -538,7 +599,7 @@ func TestProxySettingsWrittenInProxyMode(t *testing.T) { mustSetSecret(t, sec, tokenrefresh.AccessTokenKey, "existing-access") mustSetSecret(t, sec, tokenrefresh.RefreshTokenKey, "existing-refresh") mustSetSecret(t, sec, tokenrefresh.AccessTokenExpiresAtKey, fixedNow().Add(time.Hour).Format(time.RFC3339)) - var setEnvKey, setEnvValue, persistEnvKey, persistEnvValue string + var setEnvCalled, persistEnvCalled bool var daemonStarted bool _, err := Ensure(context.Background(), EnsureOptions{ @@ -546,11 +607,11 @@ func TestProxySettingsWrittenInProxyMode(t *testing.T) { Secrets: sec, Env: func(string) string { return "" }, SetEnv: func(key, value string) error { - setEnvKey, setEnvValue = key, value + setEnvCalled = true return nil }, PersistEnv: func(key, value string) error { - persistEnvKey, persistEnvValue = key, value + persistEnvCalled = true return nil }, StartDaemon: func(context.Context) error { @@ -567,13 +628,14 @@ func TestProxySettingsWrittenInProxyMode(t *testing.T) { } assertConfigContains(t, filepath.Join(tmp, "codex", "config.toml"), `base_url = "`+modelproxy.DefaultBaseURL+`"`, - `env_key = "`+codex.LocalProxyAPIKeyEnv+`"`, + `experimental_bearer_token = "`+codex.LocalProxyAPIKeyValue+`"`, ) - if setEnvKey != codex.LocalProxyAPIKeyEnv || setEnvValue != codex.LocalProxyAPIKeyValue { - t.Fatalf("SetEnv = (%q, %q), want (%q, %q)", setEnvKey, setEnvValue, codex.LocalProxyAPIKeyEnv, codex.LocalProxyAPIKeyValue) + assertConfigNotContains(t, filepath.Join(tmp, "codex", "config.toml"), `env_key = "`) + if setEnvCalled { + t.Fatal("SetEnv called for local proxy config") } - if persistEnvKey != codex.LocalProxyAPIKeyEnv || persistEnvValue != codex.LocalProxyAPIKeyValue { - t.Fatalf("PersistEnv = (%q, %q), want (%q, %q)", persistEnvKey, persistEnvValue, codex.LocalProxyAPIKeyEnv, codex.LocalProxyAPIKeyValue) + if persistEnvCalled { + t.Fatal("PersistEnv called for local proxy config") } } diff --git a/internal/modelaccess/daemon.go b/internal/modelaccess/daemon.go index 42ace94..d1f2e1d 100644 --- a/internal/modelaccess/daemon.go +++ b/internal/modelaccess/daemon.go @@ -14,7 +14,6 @@ import ( "github.com/agentserver/agentserver-pkg/internal/modelproxy" "github.com/agentserver/agentserver-pkg/internal/modelserver" "github.com/agentserver/agentserver-pkg/internal/oauth" - "github.com/agentserver/agentserver-pkg/internal/process" "github.com/agentserver/agentserver-pkg/internal/secrets" "github.com/agentserver/agentserver-pkg/internal/tokenrefresh" ) @@ -149,10 +148,11 @@ func EnsureDaemon(ctx context.Context, opts EnsureDaemonOptions) error { return err } cmd := newDaemonCommand(opts.ExePath, "model-proxy-daemon") - cmd.Stdin = nil - cmd.Stdout = nil - cmd.Stderr = nil - process.HideWindow(cmd) + cleanup, err := configureDaemonProcess(cmd) + if err != nil { + return err + } + defer cleanup() startProcess := opts.StartProcess if startProcess == nil { startProcess = startDetachedProcess diff --git a/internal/modelaccess/daemon_process_unix.go b/internal/modelaccess/daemon_process_unix.go new file mode 100644 index 0000000..f64b86e --- /dev/null +++ b/internal/modelaccess/daemon_process_unix.go @@ -0,0 +1,21 @@ +//go:build !windows + +package modelaccess + +import ( + "os" + "os/exec" + "syscall" +) + +func configureDaemonProcess(cmd *exec.Cmd) (func(), error) { + devNull, err := os.OpenFile(os.DevNull, os.O_RDWR, 0) + if err != nil { + return nil, err + } + cmd.Stdin = devNull + cmd.Stdout = devNull + cmd.Stderr = devNull + cmd.SysProcAttr = &syscall.SysProcAttr{Setsid: true} + return func() { _ = devNull.Close() }, nil +} diff --git a/internal/modelaccess/daemon_process_windows.go b/internal/modelaccess/daemon_process_windows.go new file mode 100644 index 0000000..005332a --- /dev/null +++ b/internal/modelaccess/daemon_process_windows.go @@ -0,0 +1,17 @@ +//go:build windows + +package modelaccess + +import ( + "os/exec" + + "github.com/agentserver/agentserver-pkg/internal/process" +) + +func configureDaemonProcess(cmd *exec.Cmd) (func(), error) { + cmd.Stdin = nil + cmd.Stdout = nil + cmd.Stderr = nil + process.HideWindow(cmd) + return func() {}, nil +} diff --git a/internal/modelaccess/daemon_unix_test.go b/internal/modelaccess/daemon_unix_test.go new file mode 100644 index 0000000..8497fd2 --- /dev/null +++ b/internal/modelaccess/daemon_unix_test.go @@ -0,0 +1,45 @@ +//go:build !windows + +package modelaccess + +import ( + "context" + "os" + "os/exec" + "testing" +) + +func TestEnsureDaemonDetachesProxyDaemonOnUnix(t *testing.T) { + healthy := false + err := EnsureDaemon(context.Background(), EnsureDaemonOptions{ + ExePath: "/opt/agentserver/agentserver", + ProxyBaseURL: "http://127.0.0.1:1", + HealthCheck: func(context.Context, string) bool { + return healthy + }, + StartProcess: func(cmd *exec.Cmd) error { + if cmd.SysProcAttr == nil || !cmd.SysProcAttr.Setsid { + t.Fatalf("SysProcAttr=%#v, want Setsid", cmd.SysProcAttr) + } + assertDevNullFile(t, "Stdin", cmd.Stdin) + assertDevNullFile(t, "Stdout", cmd.Stdout) + assertDevNullFile(t, "Stderr", cmd.Stderr) + healthy = true + return nil + }, + }) + if err != nil { + t.Fatalf("EnsureDaemon returned %v", err) + } +} + +func assertDevNullFile(t *testing.T, name string, got any) { + t.Helper() + f, ok := got.(*os.File) + if !ok { + t.Fatalf("%s=%T, want *os.File", name, got) + } + if f.Name() != os.DevNull { + t.Fatalf("%s file=%q, want %q", name, f.Name(), os.DevNull) + } +} diff --git a/internal/modelserver/device.go b/internal/modelserver/device.go index 37ad2b6..7de411a 100644 --- a/internal/modelserver/device.go +++ b/internal/modelserver/device.go @@ -13,8 +13,8 @@ func DeviceConfig(endpoint string) oauth.Config { } return oauth.Config{ Endpoint: strings.TrimRight(strings.TrimSpace(endpoint), "/"), - AuthPath: "/oauth2/device/auth", - TokenPath: cfg.TokenPath, + AuthPath: "/oauth/device/code", + TokenPath: "/oauth/device/token", ClientID: cfg.ClientID, Scope: cfg.Scope, } diff --git a/internal/modelserver/device_test.go b/internal/modelserver/device_test.go index 9eed5a7..6527aed 100644 --- a/internal/modelserver/device_test.go +++ b/internal/modelserver/device_test.go @@ -2,15 +2,15 @@ package modelserver import "testing" -func TestDeviceConfigUsesHydraNativeDeviceFlow(t *testing.T) { +func TestDeviceConfigUsesModelserverDeviceWrapper(t *testing.T) { got := DeviceConfig("") if got.Endpoint != "https://codeapi.cs.ac.cn" { t.Fatalf("Endpoint=%q", got.Endpoint) } - if got.AuthPath != "/oauth2/device/auth" { + if got.AuthPath != "/oauth/device/code" { t.Fatalf("AuthPath=%q", got.AuthPath) } - if got.TokenPath != "/oauth2/token" { + if got.TokenPath != "/oauth/device/token" { t.Fatalf("TokenPath=%q", got.TokenPath) } if got.ClientID != "5321f7e6-3d79-4ac9-a742-04809dbf9025" { @@ -26,7 +26,10 @@ func TestDeviceConfigAllowsEndpointOverride(t *testing.T) { if got.Endpoint != "https://codeapi.test" { t.Fatalf("Endpoint=%q", got.Endpoint) } - if got.AuthURL() != "https://codeapi.test/oauth2/device/auth" { + if got.AuthURL() != "https://codeapi.test/oauth/device/code" { t.Fatalf("AuthURL=%q", got.AuthURL()) } + if got.TokenURL() != "https://codeapi.test/oauth/device/token" { + t.Fatalf("TokenURL=%q", got.TokenURL()) + } } diff --git a/internal/ui/orchestrator_real_test.go b/internal/ui/orchestrator_real_test.go index b06fa23..c41958e 100644 --- a/internal/ui/orchestrator_real_test.go +++ b/internal/ui/orchestrator_real_test.go @@ -1278,8 +1278,8 @@ func TestConfigureCodexDesktopWritesSharedConfigOnly(t *testing.T) { if !strings.Contains(string(b), `base_url = "`+modelproxy.DefaultBaseURL+`"`) { t.Fatalf("config missing local proxy base_url:\n%s", b) } - if !strings.Contains(string(b), `env_key = "`+codex.LocalProxyAPIKeyEnv+`"`) { - t.Fatalf("config missing local proxy env_key:\n%s", b) + if !strings.Contains(string(b), `experimental_bearer_token = "`+codex.LocalProxyAPIKeyValue+`"`) { + t.Fatalf("config missing local proxy bearer token:\n%s", b) } if got := os.Getenv(codex.LocalProxyAPIKeyEnv); got != codex.LocalProxyAPIKeyValue { t.Fatalf("%s=%q, want stable local key", codex.LocalProxyAPIKeyEnv, got) From 7d2dfabb8669e565e2cb597c0b7426c1eb38fac1 Mon Sep 17 00:00:00 2001 From: Zishu Yu Date: Sun, 14 Jun 2026 04:18:02 +0800 Subject: [PATCH 39/40] fix(headless): secure local proxy access --- cmd/token-refresher/main_test.go | 8 ++- internal/headless/slave.go | 7 +++ internal/headless/slave_test.go | 58 +++++++++++++++++++++ internal/modelaccess/daemon_test.go | 8 ++- internal/modelproxy/proxy.go | 15 ++++++ internal/modelproxy/proxy_test.go | 81 ++++++++++++++++++++++++++++- 6 files changed, 173 insertions(+), 4 deletions(-) diff --git a/cmd/token-refresher/main_test.go b/cmd/token-refresher/main_test.go index aa74213..f4e14fc 100644 --- a/cmd/token-refresher/main_test.go +++ b/cmd/token-refresher/main_test.go @@ -10,6 +10,7 @@ import ( "testing" "time" + "github.com/agentserver/agentserver-pkg/internal/codex" "github.com/agentserver/agentserver-pkg/internal/modelproxy" "github.com/agentserver/agentserver-pkg/internal/oauth" "github.com/agentserver/agentserver-pkg/internal/secrets" @@ -58,7 +59,12 @@ func TestRunWithDepsServesLocalModelProxy(t *testing.T) { }() waitForProxyHealth(t, addr) - resp, err := http.Get("http://" + addr + "/v1/models") + req, err := http.NewRequest(http.MethodGet, "http://"+addr+"/v1/models", nil) + if err != nil { + t.Fatal(err) + } + req.Header.Set("Authorization", "Bearer "+codex.LocalProxyAPIKeyValue) + resp, err := http.DefaultClient.Do(req) if err != nil { t.Fatalf("proxy request: %v", err) } diff --git a/internal/headless/slave.go b/internal/headless/slave.go index ba58821..fbbcdc8 100644 --- a/internal/headless/slave.go +++ b/internal/headless/slave.go @@ -83,6 +83,13 @@ func RunSlave(ctx context.Context, opts SlaveOptions) error { name = existing.Name } else { name = defaultHeadlessSlaveName + existingSlaves, err := registry.List() + if err != nil { + return err + } + if len(existingSlaves) > 0 { + name = filepath.Base(canonicalFolder) + } if opts.NamePrompt != nil { prompted, err := opts.NamePrompt(name) if err != nil { diff --git a/internal/headless/slave_test.go b/internal/headless/slave_test.go index 7822716..c0f9c72 100644 --- a/internal/headless/slave_test.go +++ b/internal/headless/slave_test.go @@ -114,6 +114,64 @@ func TestRunSlaveDefaultsNewEntryToCommanderName(t *testing.T) { } } +func TestRunSlaveDefaultsSecondNewEntryToFolderName(t *testing.T) { + temp := t.TempDir() + repoA := filepath.Join(temp, "repo-a") + repoB := filepath.Join(temp, "repo-b") + if err := os.Mkdir(repoA, 0o755); err != nil { + t.Fatal(err) + } + if err := os.Mkdir(repoB, 0o755); err != nil { + t.Fatal(err) + } + p := testSlavePaths(temp) + + firstRunner := &fakeSlaveProcessRunner{} + if err := RunSlave(context.Background(), SlaveOptions{ + Paths: p, + Package: Package{SlaveAgent: filepath.Join(temp, "pkg", "slave-agent")}, + WorkDir: repoA, + ComputerName: "host", + NamePrompt: func(defaultName string) (string, error) { + if defaultName != "星池指挥官" { + t.Fatalf("first defaultName=%q, want 星池指挥官", defaultName) + } + return "", nil + }, + Runner: firstRunner, + Stdout: ioDiscard{}, + }); err != nil { + t.Fatalf("first RunSlave: %v", err) + } + + secondRunner := &fakeSlaveProcessRunner{} + if err := RunSlave(context.Background(), SlaveOptions{ + Paths: p, + Package: Package{SlaveAgent: filepath.Join(temp, "pkg", "slave-agent")}, + WorkDir: repoB, + ComputerName: "host", + NamePrompt: func(defaultName string) (string, error) { + if defaultName != "repo-b" { + t.Fatalf("second defaultName=%q, want repo-b", defaultName) + } + return "", nil + }, + Runner: secondRunner, + Stdout: ioDiscard{}, + }); err != nil { + t.Fatalf("second RunSlave: %v", err) + } + + configBody, err := os.ReadFile(secondRunner.request.ConfigPath) + if err != nil { + t.Fatal(err) + } + config := string(configBody) + if !strings.Contains(config, "display_name: host-repo-b") { + t.Fatalf("config missing second default display name:\n%s", config) + } +} + func TestRunSlaveReusesExistingEntryWithoutPrompt(t *testing.T) { temp := t.TempDir() repo := filepath.Join(temp, "repo") diff --git a/internal/modelaccess/daemon_test.go b/internal/modelaccess/daemon_test.go index 0b46945..525ac6f 100644 --- a/internal/modelaccess/daemon_test.go +++ b/internal/modelaccess/daemon_test.go @@ -11,6 +11,7 @@ import ( "testing" "time" + "github.com/agentserver/agentserver-pkg/internal/codex" "github.com/agentserver/agentserver-pkg/internal/modelproxy" "github.com/agentserver/agentserver-pkg/internal/oauth" "github.com/agentserver/agentserver-pkg/internal/secrets" @@ -55,7 +56,12 @@ func TestRunDaemonServesLocalModelProxyAndKeepsServingWithoutRefreshToken(t *tes t.Fatalf("AcquireDaemonLock err=%v, want %v", err, tokenrefresh.ErrDaemonAlreadyRunning) } - resp, err := http.Get("http://" + addr + "/v1/models") + req, err := http.NewRequest(http.MethodGet, "http://"+addr+"/v1/models", nil) + if err != nil { + t.Fatal(err) + } + req.Header.Set("Authorization", "Bearer "+codex.LocalProxyAPIKeyValue) + resp, err := http.DefaultClient.Do(req) if err != nil { t.Fatalf("proxy request: %v", err) } diff --git a/internal/modelproxy/proxy.go b/internal/modelproxy/proxy.go index eb29854..fb5276f 100644 --- a/internal/modelproxy/proxy.go +++ b/internal/modelproxy/proxy.go @@ -3,14 +3,17 @@ package modelproxy import ( "context" + "crypto/subtle" "errors" "fmt" "net" "net/http" "net/http/httputil" "net/url" + "strings" "time" + "github.com/agentserver/agentserver-pkg/internal/codex" "github.com/agentserver/agentserver-pkg/internal/secrets" "github.com/agentserver/agentserver-pkg/internal/tokenrefresh" ) @@ -68,6 +71,10 @@ func NewHandler(opts Options) (http.Handler, error) { w.WriteHeader(http.StatusNoContent) return } + if !validLocalBearer(r.Header.Get("Authorization")) { + http.Error(w, "local model proxy authorization required", http.StatusUnauthorized) + return + } token, err := opts.Secrets.Get(tokenrefresh.AccessTokenKey) if err != nil || token == "" { http.Error(w, "modelserver login required", http.StatusUnauthorized) @@ -80,6 +87,14 @@ func NewHandler(opts Options) (http.Handler, error) { }), nil } +func validLocalBearer(auth string) bool { + parts := strings.Fields(auth) + if len(parts) != 2 || !strings.EqualFold(parts[0], "Bearer") { + return false + } + return subtle.ConstantTimeCompare([]byte(parts[1]), []byte(codex.LocalProxyAPIKeyValue)) == 1 +} + func ListenAndServe(ctx context.Context, opts ServerOptions) error { if ctx == nil { ctx = context.Background() diff --git a/internal/modelproxy/proxy_test.go b/internal/modelproxy/proxy_test.go index 76f235c..8bbd1a8 100644 --- a/internal/modelproxy/proxy_test.go +++ b/internal/modelproxy/proxy_test.go @@ -6,6 +6,7 @@ import ( "net/http/httptest" "testing" + "github.com/agentserver/agentserver-pkg/internal/codex" "github.com/agentserver/agentserver-pkg/internal/secrets" "github.com/agentserver/agentserver-pkg/internal/tokenrefresh" ) @@ -63,7 +64,13 @@ func TestProxyLoadsLatestAccessTokenForEveryRequest(t *testing.T) { proxy := httptest.NewServer(handler) defer proxy.Close() - resp, err := http.Post(proxy.URL+"/v1/responses?trace=1", "text/plain", http.NoBody) + req, err := http.NewRequest(http.MethodPost, proxy.URL+"/v1/responses?trace=1", http.NoBody) + if err != nil { + t.Fatal(err) + } + req.Header.Set("Authorization", "Bearer "+codex.LocalProxyAPIKeyValue) + req.Header.Set("Content-Type", "text/plain") + resp, err := http.DefaultClient.Do(req) if err != nil { t.Fatalf("first request: %v", err) } @@ -78,7 +85,12 @@ func TestProxyLoadsLatestAccessTokenForEveryRequest(t *testing.T) { if err := sec.Set(tokenrefresh.AccessTokenKey, "access-2"); err != nil { t.Fatal(err) } - resp, err = http.Get(proxy.URL + "/v1/models") + req, err = http.NewRequest(http.MethodGet, proxy.URL+"/v1/models", nil) + if err != nil { + t.Fatal(err) + } + req.Header.Set("Authorization", "Bearer "+codex.LocalProxyAPIKeyValue) + resp, err = http.DefaultClient.Do(req) if err != nil { t.Fatalf("second request: %v", err) } @@ -101,6 +113,70 @@ func TestProxyLoadsLatestAccessTokenForEveryRequest(t *testing.T) { } } +func TestProxyRequiresLocalBearerToken(t *testing.T) { + sec := newTestSecrets() + if err := sec.Set(tokenrefresh.AccessTokenKey, "access"); err != nil { + t.Fatal(err) + } + var upstreamCalled bool + upstream := httptest.NewServer(http.HandlerFunc(func(w http.ResponseWriter, r *http.Request) { + upstreamCalled = true + })) + defer upstream.Close() + + handler, err := NewHandler(Options{ + Secrets: sec, + UpstreamBaseURL: upstream.URL + "/v1", + }) + if err != nil { + t.Fatalf("NewHandler: %v", err) + } + + for _, tt := range []struct { + name string + auth string + }{ + {name: "missing"}, + {name: "wrong token", auth: "Bearer wrong"}, + {name: "wrong scheme", auth: "Basic " + codex.LocalProxyAPIKeyValue}, + } { + t.Run(tt.name, func(t *testing.T) { + upstreamCalled = false + req := httptest.NewRequest(http.MethodGet, "/v1/models", nil) + if tt.auth != "" { + req.Header.Set("Authorization", tt.auth) + } + rec := httptest.NewRecorder() + handler.ServeHTTP(rec, req) + + if rec.Code != http.StatusUnauthorized { + t.Fatalf("status = %d, want %d", rec.Code, http.StatusUnauthorized) + } + if upstreamCalled { + t.Fatal("upstream should not be called without valid local bearer") + } + }) + } +} + +func TestProxyHealthDoesNotRequireLocalBearerToken(t *testing.T) { + handler, err := NewHandler(Options{ + Secrets: newTestSecrets(), + UpstreamBaseURL: "https://upstream.test/v1", + }) + if err != nil { + t.Fatalf("NewHandler: %v", err) + } + + req := httptest.NewRequest(http.MethodGet, HealthPath, nil) + rec := httptest.NewRecorder() + handler.ServeHTTP(rec, req) + + if rec.Code != http.StatusNoContent { + t.Fatalf("status = %d, want %d", rec.Code, http.StatusNoContent) + } +} + func TestProxyReturnsUnauthorizedWhenAccessTokenMissing(t *testing.T) { upstreamCalled := false upstream := httptest.NewServer(http.HandlerFunc(func(w http.ResponseWriter, r *http.Request) { @@ -117,6 +193,7 @@ func TestProxyReturnsUnauthorizedWhenAccessTokenMissing(t *testing.T) { } req := httptest.NewRequest(http.MethodGet, "/v1/models", nil) + req.Header.Set("Authorization", "Bearer "+codex.LocalProxyAPIKeyValue) rec := httptest.NewRecorder() handler.ServeHTTP(rec, req) From c2b04b363a33dfc0127ce933431cb780285e55be Mon Sep 17 00:00:00 2001 From: Zishu Yu Date: Sun, 14 Jun 2026 12:18:19 +0800 Subject: [PATCH 40/40] fix(headless): address linux proxy review --- Makefile | 2 +- cmd/agentctl/cmd_test_subcommands.go | 8 +- cmd/agentctl/cmd_test_subcommands_test.go | 2 +- cmd/agentserver/main.go | 67 +++++- cmd/agentserver/main_test.go | 76 +++++++ cmd/launcher/main.go | 14 +- cmd/launcher/main_test.go | 4 +- cmd/open-folder/main.go | 8 +- cmd/open-folder/main_test.go | 4 +- cmd/token-refresher/main.go | 15 +- cmd/token-refresher/main_test.go | 4 +- ...2026-06-13-linux-headless-server-design.md | 6 +- internal/agentserver/client.go | 9 + internal/agentserver/client_test.go | 17 ++ internal/agentserver/oauth.go | 4 +- internal/codex/config.go | 102 ++++++++-- internal/codex/config_test.go | 110 +++++++++- internal/headless/driver.go | 106 ++++++++-- internal/headless/driver_test.go | 190 +++++++++++++++++- internal/headless/runtime.go | 79 +++++++- internal/headless/runtime_test.go | 73 +++++++ internal/headless/slave.go | 3 + internal/headless/slave_test.go | 51 +++++ internal/launchprep/prepare.go | 2 +- internal/modelaccess/access.go | 22 +- internal/modelaccess/access_test.go | 94 ++++++++- internal/modelaccess/daemon.go | 14 +- internal/modelaccess/daemon_process_unix.go | 26 ++- .../modelaccess/daemon_process_windows.go | 28 ++- internal/modelaccess/daemon_test.go | 45 +++-- internal/modelaccess/daemon_unix_test.go | 16 +- internal/modelaccess/proxy_token.go | 81 ++++++++ internal/modelproxy/proxy.go | 77 +++++-- internal/modelproxy/proxy_test.go | 117 +++++++++-- internal/oauth/devicecode.go | 15 +- internal/oauth/devicecode_test.go | 23 +++ internal/slave/file_lock_unix.go | 43 ++++ internal/slave/file_lock_windows.go | 45 +++++ internal/slave/process_liveness.go | 6 + internal/slave/registry.go | 58 ++++-- internal/slave/registry_test.go | 186 ++++++++++++++++- internal/ui/orchestrator_real.go | 8 +- internal/ui/orchestrator_real_test.go | 8 +- internal/vscode/linux_package_test.go | 58 ++++++ .../linux/codex-manifest-linux-amd64.json | 3 +- .../linux/codex-manifest-linux-arm64.json | 3 +- scripts/linux-package-common.sh | 2 +- scripts/package-linux.sh | 15 +- 48 files changed, 1749 insertions(+), 200 deletions(-) create mode 100644 internal/modelaccess/proxy_token.go create mode 100644 internal/slave/file_lock_unix.go create mode 100644 internal/slave/file_lock_windows.go create mode 100644 internal/slave/process_liveness.go diff --git a/Makefile b/Makefile index 472947d..d59a586 100644 --- a/Makefile +++ b/Makefile @@ -49,7 +49,7 @@ cross-linux: @mkdir -p $(DIST)/linux/amd64 $(DIST)/linux/arm64 @for arch in amd64 arm64; do \ echo "==> cross-building agentserver (linux/$$arch)"; \ - GOOS=linux GOARCH=$$arch \ + CGO_ENABLED=0 GOOS=linux GOARCH=$$arch \ $(GO) build $(GOFLAGS) -ldflags="$(LDFLAGS)" \ -o $(DIST)/linux/$$arch/agentserver ./cmd/agentserver ; \ done diff --git a/cmd/agentctl/cmd_test_subcommands.go b/cmd/agentctl/cmd_test_subcommands.go index 7b07139..0b2b3be 100644 --- a/cmd/agentctl/cmd_test_subcommands.go +++ b/cmd/agentctl/cmd_test_subcommands.go @@ -126,7 +126,7 @@ func runTestConfigure() { die(err) } fmt.Printf("wrote codex config: %s\n", p.CodexConfigFile) - fmt.Printf("setx %s=%s (HKCU\\Environment)\n", codex.LocalProxyAPIKeyEnv, codex.LocalProxyAPIKeyValue) + fmt.Printf("setx %s=%s (HKCU\\Environment)\n", codex.LocalProxyAPIKeyEnv, codex.LegacyLocalProxyAPIKeyValue) // .vsix sits next to the running agentctl.exe exeDir, _ := os.Executable() @@ -254,13 +254,13 @@ func openTestFolder(ctx context.Context, s *state.State, p paths.Paths, folder s } func configureTestCodex(p paths.Paths) error { - if err := codex.UpdateConfig(p.CodexConfigFile, codex.ModelserverProxySettings(modelproxy.DefaultBaseURL)); err != nil { + if err := codex.UpdateConfig(p.CodexConfigFile, codex.ModelserverProxySettings(modelproxy.DefaultBaseURL, codex.LegacyLocalProxyAPIKeyValue)); err != nil { return err } - if err := env.PersistUserEnv(codex.LocalProxyAPIKeyEnv, codex.LocalProxyAPIKeyValue); err != nil { + if err := env.PersistUserEnv(codex.LocalProxyAPIKeyEnv, codex.LegacyLocalProxyAPIKeyValue); err != nil { return err } - return os.Setenv(codex.LocalProxyAPIKeyEnv, codex.LocalProxyAPIKeyValue) + return os.Setenv(codex.LocalProxyAPIKeyEnv, codex.LegacyLocalProxyAPIKeyValue) } func startTestVSCode(codeExe string, args []string) (int, error) { diff --git a/cmd/agentctl/cmd_test_subcommands_test.go b/cmd/agentctl/cmd_test_subcommands_test.go index e290792..6b7796d 100644 --- a/cmd/agentctl/cmd_test_subcommands_test.go +++ b/cmd/agentctl/cmd_test_subcommands_test.go @@ -53,7 +53,7 @@ func TestOpenTestFolderCodexDesktopUsesDeepLinkAndWritesConfig(t *testing.T) { if !strings.Contains(string(b), `base_url = "`+modelproxy.DefaultBaseURL+`"`) { t.Fatalf("config missing local proxy base_url:\n%s", b) } - if !strings.Contains(string(b), `experimental_bearer_token = "`+codex.LocalProxyAPIKeyValue+`"`) { + if !strings.Contains(string(b), `experimental_bearer_token = "`+codex.LegacyLocalProxyAPIKeyValue+`"`) { t.Fatalf("config missing local proxy bearer token:\n%s", b) } } diff --git a/cmd/agentserver/main.go b/cmd/agentserver/main.go index 6e1ebba..31452be 100644 --- a/cmd/agentserver/main.go +++ b/cmd/agentserver/main.go @@ -8,8 +8,10 @@ import ( "io" "log" "os" + "os/signal" "path/filepath" "strings" + "syscall" "github.com/agentserver/agentserver-pkg/internal/headless" "github.com/agentserver/agentserver-pkg/internal/modelaccess" @@ -18,6 +20,7 @@ import ( "github.com/agentserver/agentserver-pkg/internal/paths" "github.com/agentserver/agentserver-pkg/internal/secrets" "github.com/agentserver/agentserver-pkg/internal/terminalauth" + "golang.org/x/term" ) type app struct { @@ -32,11 +35,17 @@ type app struct { } func main() { - if err := newApp().run(context.Background(), os.Args[1:]); err != nil { + ctx, stop := commandContext(context.Background()) + defer stop() + if err := newApp().run(ctx, os.Args[1:]); err != nil { log.Fatalf("agentserver: %v", err) } } +func commandContext(parent context.Context) (context.Context, context.CancelFunc) { + return signal.NotifyContext(parent, os.Interrupt, syscall.SIGTERM) +} + func newApp() app { p, err := paths.Default() if err != nil { @@ -54,12 +63,19 @@ func newApp() app { } pkg := headless.PackagePaths(exe) sec := secrets.New(p.SecretsFile) + proxyTokenPath := modelaccess.DefaultLocalProxyTokenPath(p.InstallRoot) + proxyLogPath := filepath.Join(p.InstallRoot, "logs", "model-proxy-daemon.log") cachedCodex := "" ensureModelAccess := func(out io.Writer) func(context.Context) error { return func(ctx context.Context) error { + proxyToken, err := modelaccess.EnsureLocalProxyToken(proxyTokenPath) + if err != nil { + return err + } if _, err := modelaccess.Ensure(ctx, modelaccess.EnsureOptions{ CodexConfigPath: p.CodexConfigFile, Secrets: sec, + LocalProxyToken: proxyToken, PrintChallenge: func(title string, ch oauth.DeviceCodeChallenge) { terminalauth.PrintChallenge(out, title, ch, terminalauth.DefaultQR) }, @@ -67,6 +83,7 @@ func newApp() app { return modelaccess.EnsureDaemon(ctx, modelaccess.EnsureDaemonOptions{ ExePath: pkg.AgentserverExe, ProxyBaseURL: "http://127.0.0.1:53452", + LogPath: proxyLogPath, }) }, }); err != nil { @@ -91,7 +108,10 @@ func newApp() app { return nil }, runSlave: func(ctx context.Context) error { - wd, _ := os.Getwd() + wd, err := os.Getwd() + if err != nil { + return fmt.Errorf("get current workdir: %w", err) + } return headless.RunSlave(ctx, headless.SlaveOptions{ Paths: p, Package: pkg, @@ -121,7 +141,10 @@ func newApp() app { }) }, serveDriverMCP: func(ctx context.Context) error { - wd, _ := os.Getwd() + wd, err := os.Getwd() + if err != nil { + return fmt.Errorf("get current workdir: %w", err) + } return headless.ServeDriverMCP(ctx, headless.DriverMCPOptions{ Paths: p, Package: pkg, @@ -130,15 +153,36 @@ func newApp() app { }) }, runDaemon: func(ctx context.Context) error { + proxyToken, err := modelaccess.EnsureLocalProxyToken(proxyTokenPath) + if err != nil { + return err + } return modelaccess.RunDaemon(ctx, modelaccess.DaemonOptions{ - Secrets: sec, - OAuth: modelserver.OAuthConfig(), - LockPath: filepath.Join(p.InstallRoot, "token-refresher.lock"), + Secrets: sec, + OAuth: modelserver.OAuthConfig(), + LocalProxyToken: proxyToken, + LockPath: filepath.Join(p.InstallRoot, "token-refresher.lock"), + Logf: daemonLogf(proxyLogPath), }) }, } } +func daemonLogf(path string) func(string, ...any) { + return func(format string, args ...any) { + if err := os.MkdirAll(filepath.Dir(path), 0o700); err != nil { + return + } + f, err := os.OpenFile(path, os.O_CREATE|os.O_WRONLY|os.O_APPEND, 0o600) + if err != nil { + return + } + defer f.Close() + logger := log.New(f, "", log.LstdFlags) + logger.Printf(format, args...) + } +} + func (a app) run(ctx context.Context, args []string) error { cmd := "" if len(args) > 0 { @@ -181,7 +225,16 @@ func (a app) run(ctx context.Context, args []string) error { } func promptName(defaultName string) (string, error) { - return promptNameWithIO(os.Stdin, os.Stdout, defaultName) + return promptNameWithTerminal(os.Stdin, os.Stdout, defaultName, func() bool { + return term.IsTerminal(int(os.Stdin.Fd())) + }) +} + +func promptNameWithTerminal(r io.Reader, w io.Writer, defaultName string, isTerminal func() bool) (string, error) { + if isTerminal == nil || !isTerminal() { + return defaultName, nil + } + return promptNameWithIO(r, w, defaultName) } func promptNameWithIO(r io.Reader, w io.Writer, defaultName string) (string, error) { diff --git a/cmd/agentserver/main_test.go b/cmd/agentserver/main_test.go index 1578d6f..9425367 100644 --- a/cmd/agentserver/main_test.go +++ b/cmd/agentserver/main_test.go @@ -3,8 +3,11 @@ package main import ( "bytes" "context" + "os" + "path/filepath" "strings" "testing" + "time" ) func TestRunCommandEnsuresAccessForUserCommands(t *testing.T) { @@ -173,6 +176,79 @@ func TestRunCommandRejectsUnknownCommand(t *testing.T) { } } +func TestCommandContextCancelsOnInterrupt(t *testing.T) { + ctx, stop := commandContext(context.Background()) + defer stop() + + p, err := os.FindProcess(os.Getpid()) + if err != nil { + t.Fatal(err) + } + if err := p.Signal(os.Interrupt); err != nil { + t.Fatal(err) + } + + select { + case <-ctx.Done(): + case <-time.After(time.Second): + t.Fatal("context was not canceled after interrupt") + } +} + +func TestPromptNameSkipsPromptWhenInputIsNotTerminal(t *testing.T) { + var out bytes.Buffer + got, err := promptNameWithTerminal(strings.NewReader("custom\n"), &out, "default", func() bool { + return false + }) + if err != nil { + t.Fatal(err) + } + if got != "default" { + t.Fatalf("name=%q, want default", got) + } + if out.Len() != 0 { + t.Fatalf("prompt output=%q, want empty", out.String()) + } +} + +func TestPromptNameReadsTerminalInput(t *testing.T) { + var out bytes.Buffer + got, err := promptNameWithTerminal(strings.NewReader("custom\n"), &out, "default", func() bool { + return true + }) + if err != nil { + t.Fatal(err) + } + if got != "custom" { + t.Fatalf("name=%q, want custom", got) + } + if !strings.Contains(out.String(), "Slave name [default]: ") { + t.Fatalf("prompt output=%q", out.String()) + } +} + +func TestDaemonLogfAppendsToPrivateLogFile(t *testing.T) { + path := filepath.Join(t.TempDir(), "logs", "model-proxy-daemon.log") + logf := daemonLogf(path) + + logf("refresh failed: %s", "boom") + + body, err := os.ReadFile(path) + if err != nil { + t.Fatal(err) + } + if !strings.Contains(string(body), "refresh failed: boom") { + t.Fatalf("log body=%q", body) + } + info, err := os.Stat(path) + if err != nil { + t.Fatal(err) + } + if got := info.Mode().Perm(); got != 0o600 { + t.Fatalf("log mode=%#o, want 0600", got) + } +} + func TestPromptNameWithIO(t *testing.T) { for _, tc := range []struct { name string diff --git a/cmd/launcher/main.go b/cmd/launcher/main.go index e327896..4dda625 100644 --- a/cmd/launcher/main.go +++ b/cmd/launcher/main.go @@ -788,11 +788,11 @@ func launchCompletedFrontend(ctx context.Context, s *state.State, p paths.Paths, } func launchCompletedCodexDesktop(ctx context.Context, s *state.State, p paths.Paths, sec secrets.Store, installDir string, tokenRefresherExe string, opener codexdesktop.Opener) error { - if err := codex.UpdateConfig(p.CodexConfigFile, codex.ModelserverProxySettings(modelproxy.DefaultBaseURL)); err != nil { + if err := codex.UpdateConfig(p.CodexConfigFile, codex.ModelserverProxySettings(modelproxy.DefaultBaseURL, codex.LegacyLocalProxyAPIKeyValue)); err != nil { return err } - _ = env.PersistUserEnv(codex.LocalProxyAPIKeyEnv, codex.LocalProxyAPIKeyValue) - _ = os.Setenv(codex.LocalProxyAPIKeyEnv, codex.LocalProxyAPIKeyValue) + _ = env.PersistUserEnv(codex.LocalProxyAPIKeyEnv, codex.LegacyLocalProxyAPIKeyValue) + _ = os.Setenv(codex.LocalProxyAPIKeyEnv, codex.LegacyLocalProxyAPIKeyValue) if err := codexdesktop.ConfigureLocale( p.CodexDesktopGlobalStateFile, p.CodexDesktopComputerUseConfigFile, @@ -954,17 +954,17 @@ func lastN(s string, n int) string { } func execVSCode(codeExe string, p paths.Paths, folder string, sec secrets.Store, tokenRefresherExe string) error { - if err := codex.UpdateConfig(p.CodexConfigFile, codex.ModelserverProxySettings(modelproxy.DefaultBaseURL)); err != nil { + if err := codex.UpdateConfig(p.CodexConfigFile, codex.ModelserverProxySettings(modelproxy.DefaultBaseURL, codex.LegacyLocalProxyAPIKeyValue)); err != nil { return err } - _ = env.PersistUserEnv(codex.LocalProxyAPIKeyEnv, codex.LocalProxyAPIKeyValue) - _ = os.Setenv(codex.LocalProxyAPIKeyEnv, codex.LocalProxyAPIKeyValue) + _ = env.PersistUserEnv(codex.LocalProxyAPIKeyEnv, codex.LegacyLocalProxyAPIKeyValue) + _ = os.Setenv(codex.LocalProxyAPIKeyEnv, codex.LegacyLocalProxyAPIKeyValue) if tokenRefresherExe != "" { _ = tokenrefresh.StartDaemon(tokenRefresherExe) } args := vscode.LaunchArgs(p.VSCodeUserDataDir, p.VSCodeExtDir, folder) cmd := exec.Command(codeExe, args...) - cmd.Env = vscode.UpsertEnv(os.Environ(), codex.LocalProxyAPIKeyEnv, codex.LocalProxyAPIKeyValue) + cmd.Env = vscode.UpsertEnv(os.Environ(), codex.LocalProxyAPIKeyEnv, codex.LegacyLocalProxyAPIKeyValue) cmd.Stdout = os.Stdout cmd.Stderr = os.Stderr return cmd.Start() diff --git a/cmd/launcher/main_test.go b/cmd/launcher/main_test.go index 1d67d7e..e5decba 100644 --- a/cmd/launcher/main_test.go +++ b/cmd/launcher/main_test.go @@ -893,7 +893,7 @@ func TestExecVSCodeEnsuresCodexConfigBeforeLaunch(t *testing.T) { for _, want := range []string{ `model_provider = "modelserver"`, `base_url = "` + modelproxy.DefaultBaseURL + `"`, - `experimental_bearer_token = "` + codex.LocalProxyAPIKeyValue + `"`, + `experimental_bearer_token = "` + codex.LegacyLocalProxyAPIKeyValue + `"`, `[windows]`, `sandbox = "unelevated"`, } { @@ -992,7 +992,7 @@ func TestLaunchCompletedCodexDesktopWritesConfigAndOpensDeepLink(t *testing.T) { if !strings.Contains(string(b), `base_url = "`+modelproxy.DefaultBaseURL+`"`) { t.Fatalf("config missing local proxy base_url:\n%s", b) } - if !strings.Contains(string(b), `experimental_bearer_token = "`+codex.LocalProxyAPIKeyValue+`"`) { + if !strings.Contains(string(b), `experimental_bearer_token = "`+codex.LegacyLocalProxyAPIKeyValue+`"`) { t.Fatalf("config missing local proxy bearer token:\n%s", b) } } diff --git a/cmd/open-folder/main.go b/cmd/open-folder/main.go index c49227f..19c8769 100644 --- a/cmd/open-folder/main.go +++ b/cmd/open-folder/main.go @@ -127,18 +127,18 @@ func openFolder(ctx context.Context, codeExe string, p paths.Paths, folder strin } cmd := exec.Command(codeExe, vscode.LaunchArgs(p.VSCodeUserDataDir, p.VSCodeExtDir, folder)...) - cmd.Env = vscode.UpsertEnv(os.Environ(), codex.LocalProxyAPIKeyEnv, codex.LocalProxyAPIKeyValue) + cmd.Env = vscode.UpsertEnv(os.Environ(), codex.LocalProxyAPIKeyEnv, codex.LegacyLocalProxyAPIKeyValue) cmd.Stdout = os.Stdout cmd.Stderr = os.Stderr return cmd.Start() } func openFolderCodexDesktop(ctx context.Context, p paths.Paths, folder string, sec secrets.Store, tokenRefresherExe string, opener codexdesktop.Opener) error { - if err := codex.UpdateConfig(p.CodexConfigFile, codex.ModelserverProxySettings(modelproxy.DefaultBaseURL)); err != nil { + if err := codex.UpdateConfig(p.CodexConfigFile, codex.ModelserverProxySettings(modelproxy.DefaultBaseURL, codex.LegacyLocalProxyAPIKeyValue)); err != nil { return err } - _ = env.PersistUserEnv(codex.LocalProxyAPIKeyEnv, codex.LocalProxyAPIKeyValue) - _ = os.Setenv(codex.LocalProxyAPIKeyEnv, codex.LocalProxyAPIKeyValue) + _ = env.PersistUserEnv(codex.LocalProxyAPIKeyEnv, codex.LegacyLocalProxyAPIKeyValue) + _ = os.Setenv(codex.LocalProxyAPIKeyEnv, codex.LegacyLocalProxyAPIKeyValue) if err := codexdesktop.ConfigureLocale( p.CodexDesktopGlobalStateFile, p.CodexDesktopComputerUseConfigFile, diff --git a/cmd/open-folder/main_test.go b/cmd/open-folder/main_test.go index 50ab700..b629ec0 100644 --- a/cmd/open-folder/main_test.go +++ b/cmd/open-folder/main_test.go @@ -63,7 +63,7 @@ func TestOpenFolderMigratesVSCodeSettingsBeforeLaunch(t *testing.T) { } for _, want := range []string{ `base_url = "` + modelproxy.DefaultBaseURL + `"`, - `experimental_bearer_token = "` + codex.LocalProxyAPIKeyValue + `"`, + `experimental_bearer_token = "` + codex.LegacyLocalProxyAPIKeyValue + `"`, } { if !strings.Contains(string(b), want) { t.Fatalf("missing %q in:\n%s", want, b) @@ -94,7 +94,7 @@ func TestOpenFolderCodexDesktopUsesFolderDeepLink(t *testing.T) { } for _, want := range []string{ `base_url = "` + modelproxy.DefaultBaseURL + `"`, - `experimental_bearer_token = "` + codex.LocalProxyAPIKeyValue + `"`, + `experimental_bearer_token = "` + codex.LegacyLocalProxyAPIKeyValue + `"`, } { if !strings.Contains(string(b), want) { t.Fatalf("missing %q in:\n%s", want, b) diff --git a/cmd/token-refresher/main.go b/cmd/token-refresher/main.go index 2a87b1c..6aa5dd4 100644 --- a/cmd/token-refresher/main.go +++ b/cmd/token-refresher/main.go @@ -30,17 +30,23 @@ func run() error { if err != nil { return err } + proxyToken, err := modelaccess.EnsureLocalProxyToken(modelaccess.DefaultLocalProxyTokenPath(p.InstallRoot)) + if err != nil { + return err + } return runWithDeps(context.Background(), runDeps{ - Secrets: secrets.New(p.SecretsFile), - OAuth: modelserver.OAuthConfig(), - LockPath: filepath.Join(p.InstallRoot, "token-refresher.lock"), - Logf: log.Printf, + Secrets: secrets.New(p.SecretsFile), + OAuth: modelserver.OAuthConfig(), + LocalProxyToken: proxyToken, + LockPath: filepath.Join(p.InstallRoot, "token-refresher.lock"), + Logf: log.Printf, }) } type runDeps struct { Secrets secrets.Store OAuth oauth.AuthCodeConfig + LocalProxyToken string ProxyAddr string ProxyUpstreamBaseURL string LockPath string @@ -58,6 +64,7 @@ func runWithDeps(ctx context.Context, deps runDeps) error { return modelaccess.RunDaemon(ctx, modelaccess.DaemonOptions{ Secrets: deps.Secrets, OAuth: deps.OAuth, + LocalProxyToken: deps.LocalProxyToken, ProxyAddr: deps.ProxyAddr, ProxyUpstreamBaseURL: deps.ProxyUpstreamBaseURL, LockPath: deps.LockPath, diff --git a/cmd/token-refresher/main_test.go b/cmd/token-refresher/main_test.go index f4e14fc..cb992eb 100644 --- a/cmd/token-refresher/main_test.go +++ b/cmd/token-refresher/main_test.go @@ -10,7 +10,6 @@ import ( "testing" "time" - "github.com/agentserver/agentserver-pkg/internal/codex" "github.com/agentserver/agentserver-pkg/internal/modelproxy" "github.com/agentserver/agentserver-pkg/internal/oauth" "github.com/agentserver/agentserver-pkg/internal/secrets" @@ -46,6 +45,7 @@ func TestRunWithDepsServesLocalModelProxy(t *testing.T) { errCh <- runWithDeps(ctx, runDeps{ Secrets: sec, OAuth: oauth.AuthCodeConfig{ClientID: "client-x"}, + LocalProxyToken: "random-local-token", ProxyAddr: addr, ProxyUpstreamBaseURL: upstream.URL + "/v1", Refresh: func(context.Context, oauth.AuthCodeConfig, string) (oauth.Token, error) { @@ -63,7 +63,7 @@ func TestRunWithDepsServesLocalModelProxy(t *testing.T) { if err != nil { t.Fatal(err) } - req.Header.Set("Authorization", "Bearer "+codex.LocalProxyAPIKeyValue) + req.Header.Set("Authorization", "Bearer random-local-token") resp, err := http.DefaultClient.Do(req) if err != nil { t.Fatalf("proxy request: %v", err) diff --git a/docs/superpowers/specs/2026-06-13-linux-headless-server-design.md b/docs/superpowers/specs/2026-06-13-linux-headless-server-design.md index 4f75ebd..686b5e8 100644 --- a/docs/superpowers/specs/2026-06-13-linux-headless-server-design.md +++ b/docs/superpowers/specs/2026-06-13-linux-headless-server-design.md @@ -75,11 +75,13 @@ check has two modes, in this order: [model_providers.modelserver] name = "modelserver" base_url = "http://127.0.0.1:53452/v1" - env_key = "AGENTSERVER_CODEX_LOCAL_API_KEY" + experimental_bearer_token = "" wire_api = "responses" ``` -In proxy mode, the local proxy injects the latest access token from the secrets +In proxy mode, `~/.agentserver-app/proxy-token` is a per-user random bearer token +written with mode `0600`. The local proxy rejects non-health requests unless +that bearer is present, then injects the latest access token from the secrets store for every request to `https://code.ai.cs.ac.cn/v1`. The refresher uses the stored refresh token and marks reauth-required on `invalid_grant` or missing refresh token. The next user-facing command reruns code device login. diff --git a/internal/agentserver/client.go b/internal/agentserver/client.go index 1c35cdb..f41f486 100644 --- a/internal/agentserver/client.go +++ b/internal/agentserver/client.go @@ -8,6 +8,7 @@ import ( "fmt" "io" "net/http" + "net/url" "strings" "time" ) @@ -77,6 +78,14 @@ func (c *Client) RegisterAgent(ctx context.Context, oauthToken, name, typ string return out, nil } +func (c *Client) DeleteAgent(ctx context.Context, proxyToken, sandboxID string) error { + sandboxID = strings.TrimSpace(sandboxID) + if sandboxID == "" { + return fmt.Errorf("delete agent: sandbox id required") + } + return c.do(ctx, http.MethodDelete, "/api/agent/sandboxes/"+url.PathEscape(sandboxID), proxyToken, nil, nil) +} + func (c *Client) CreateWorkspaceAPIKey(ctx context.Context, token, workspaceID, name string) (WorkspaceAPIKey, error) { var wrap struct { Data WorkspaceAPIKey `json:"data"` diff --git a/internal/agentserver/client_test.go b/internal/agentserver/client_test.go index 533fc07..cb61eca 100644 --- a/internal/agentserver/client_test.go +++ b/internal/agentserver/client_test.go @@ -116,6 +116,23 @@ func TestRegisterAgentUsesOAuthToken(t *testing.T) { } } +func TestDeleteAgentUsesSandboxProxyToken(t *testing.T) { + srv := httptest.NewServer(http.HandlerFunc(func(w http.ResponseWriter, r *http.Request) { + if r.Method != http.MethodDelete || r.URL.Path != "/api/agent/sandboxes/sb-1" { + t.Fatalf("got %s %s", r.Method, r.URL.Path) + } + if r.Header.Get("Authorization") != "Bearer sandbox-proxy-token" { + t.Fatalf("Authorization=%q", r.Header.Get("Authorization")) + } + w.WriteHeader(http.StatusNoContent) + })) + defer srv.Close() + + if err := New(srv.URL).DeleteAgent(context.Background(), "sandbox-proxy-token", "sb-1"); err != nil { + t.Fatal(err) + } +} + func TestNewClientUsesBoundedHTTPTimeout(t *testing.T) { got := New("https://agent.cs.ac.cn").http.Timeout if got != 60*time.Second { diff --git a/internal/agentserver/oauth.go b/internal/agentserver/oauth.go index e2bf22e..cadc3b2 100644 --- a/internal/agentserver/oauth.go +++ b/internal/agentserver/oauth.go @@ -6,9 +6,11 @@ import ( "github.com/agentserver/agentserver-pkg/internal/oauth" ) +const DefaultEndpoint = "https://agent.cs.ac.cn" + func OAuthConfig(endpoint string) oauth.Config { if strings.TrimSpace(endpoint) == "" { - endpoint = "https://agent.cs.ac.cn" + endpoint = DefaultEndpoint } return oauth.Config{ Endpoint: strings.TrimRight(strings.TrimSpace(endpoint), "/"), diff --git a/internal/codex/config.go b/internal/codex/config.go index 728ea7a..1e58598 100644 --- a/internal/codex/config.go +++ b/internal/codex/config.go @@ -50,15 +50,17 @@ func ModelserverSettings() Settings { } const ( - LocalProxyAPIKeyEnv = "AGENTSERVER_CODEX_LOCAL_API_KEY" - LocalProxyAPIKeyValue = "agentserver-local-proxy" + LocalProxyAPIKeyEnv = "AGENTSERVER_CODEX_LOCAL_API_KEY" + // LegacyLocalProxyAPIKeyValue is retained for the Windows desktop launcher + // path. Linux headless model proxy access must use a per-user random token. + LegacyLocalProxyAPIKeyValue = "agentserver-local-proxy" ) -func ModelserverProxySettings(baseURL string) Settings { +func ModelserverProxySettings(baseURL, bearerToken string) Settings { s := ModelserverSettings() s.BaseURL = baseURL s.EnvKey = "" - s.ExperimentalBearerToken = LocalProxyAPIKeyValue + s.ExperimentalBearerToken = strings.TrimSpace(bearerToken) return s } @@ -87,7 +89,9 @@ func UpdateConfig(path string, s Settings) error { if _, err := toml.Decode(string(b), &root); err != nil { return fmt.Errorf("parse existing config.toml: %w", err) } - _ = writeConfigBackup(path, b) + if err := writeConfigBackup(path, b); err != nil { + return fmt.Errorf("backup config.toml: %w", err) + } } else if !errors.Is(err, os.ErrNotExist) { return fmt.Errorf("read config.toml: %w", err) } @@ -110,16 +114,22 @@ func UpdateConfig(path string, s Settings) error { if providers == nil { providers = map[string]any{} } - provider := map[string]any{ - "name": s.Provider, - "base_url": s.BaseURL, - "wire_api": s.WireAPI, + provider, _ := providers[s.Provider].(map[string]any) + if provider == nil { + provider = map[string]any{} } + provider["name"] = s.Provider + provider["base_url"] = s.BaseURL + provider["wire_api"] = s.WireAPI if s.EnvKey != "" { provider["env_key"] = s.EnvKey + } else { + delete(provider, "env_key") } if s.ExperimentalBearerToken != "" { provider["experimental_bearer_token"] = s.ExperimentalBearerToken + } else { + delete(provider, "experimental_bearer_token") } providers[s.Provider] = provider root["model_providers"] = providers @@ -128,7 +138,7 @@ func UpdateConfig(path string, s Settings) error { if err := toml.NewEncoder(&buf).Encode(root); err != nil { return fmt.Errorf("marshal config.toml: %w", err) } - return os.WriteFile(path, buf.Bytes(), 0o644) + return writeConfigFile(path, buf.Bytes()) } func HasModelserverDirectConfig(path string) (bool, error) { @@ -187,7 +197,9 @@ func UpdateMCPServer(path, name string, server MCPServer) error { if _, err := toml.Decode(string(b), &root); err != nil { return fmt.Errorf("parse existing config.toml: %w", err) } - _ = writeConfigBackup(path, b) + if err := writeConfigBackup(path, b); err != nil { + return fmt.Errorf("backup config.toml: %w", err) + } } else if !errors.Is(err, os.ErrNotExist) { return fmt.Errorf("read config.toml: %w", err) } @@ -196,24 +208,36 @@ func UpdateMCPServer(path, name string, server MCPServer) error { if servers == nil { servers = map[string]any{} } - entry := map[string]any{ - "command": server.Command, - "args": append([]string(nil), server.Args...), + entry, _ := servers[name].(map[string]any) + if entry == nil { + entry = map[string]any{} } + entry["command"] = server.Command + entry["args"] = append([]string(nil), server.Args...) if server.StartupTimeoutSec > 0 { entry["startup_timeout_sec"] = server.StartupTimeoutSec + } else { + delete(entry, "startup_timeout_sec") } if server.ToolTimeoutSec > 0 { entry["tool_timeout_sec"] = server.ToolTimeoutSec + } else { + delete(entry, "tool_timeout_sec") } if server.Enabled != nil { entry["enabled"] = *server.Enabled + } else { + delete(entry, "enabled") } if len(server.Env) > 0 { entry["env"] = server.Env + } else { + delete(entry, "env") } if server.Cwd != "" { entry["cwd"] = server.Cwd + } else { + delete(entry, "cwd") } servers[name] = entry root["mcp_servers"] = servers @@ -222,7 +246,7 @@ func UpdateMCPServer(path, name string, server MCPServer) error { if err := toml.NewEncoder(&buf).Encode(root); err != nil { return fmt.Errorf("marshal config.toml: %w", err) } - return os.WriteFile(path, buf.Bytes(), 0o644) + return writeConfigFile(path, buf.Bytes()) } func RemoveMCPServer(path, name string) error { @@ -247,7 +271,9 @@ func RemoveMCPServer(path, name string) error { if _, ok := servers[name]; !ok { return nil } - _ = writeConfigBackup(path, b) + if err := writeConfigBackup(path, b); err != nil { + return fmt.Errorf("backup config.toml: %w", err) + } delete(servers, name) if len(servers) == 0 { delete(root, "mcp_servers") @@ -258,7 +284,7 @@ func RemoveMCPServer(path, name string) error { if err := toml.NewEncoder(&buf).Encode(root); err != nil { return fmt.Errorf("marshal config.toml: %w", err) } - return os.WriteFile(path, buf.Bytes(), 0o644) + return writeConfigFile(path, buf.Bytes()) } func defaultString(v, fallback string) string { @@ -268,14 +294,52 @@ func defaultString(v, fallback string) string { return fallback } -func writeConfigBackup(path string, body []byte) error { +var writeConfigBackup = writeConfigBackupFile + +func writeConfigBackupFile(path string, body []byte) error { backup := fmt.Sprintf("%s.bak.%d", path, time.Now().UnixNano()) - if err := os.WriteFile(backup, body, 0o644); err != nil { + if err := writeConfigFile(backup, body); err != nil { return err } return pruneConfigBackups(path) } +func writeConfigFile(path string, body []byte) error { + if err := os.MkdirAll(filepath.Dir(path), 0o755); err != nil { + return err + } + tmp, err := os.CreateTemp(filepath.Dir(path), filepath.Base(path)+".*.tmp") + if err != nil { + return err + } + tmpPath := tmp.Name() + defer func() { + if tmpPath != "" { + _ = os.Remove(tmpPath) + } + }() + if err := tmp.Chmod(0o600); err != nil { + _ = tmp.Close() + return err + } + if _, err := tmp.Write(body); err != nil { + _ = tmp.Close() + return err + } + if err := tmp.Sync(); err != nil { + _ = tmp.Close() + return err + } + if err := tmp.Close(); err != nil { + return err + } + if err := os.Rename(tmpPath, path); err != nil { + return err + } + tmpPath = "" + return os.Chmod(path, 0o600) +} + func pruneConfigBackups(path string) error { matches, err := filepath.Glob(path + ".bak.*") if err != nil { diff --git a/internal/codex/config_test.go b/internal/codex/config_test.go index 8280d84..2ee6caf 100644 --- a/internal/codex/config_test.go +++ b/internal/codex/config_test.go @@ -1,6 +1,7 @@ package codex import ( + "errors" "os" "path/filepath" "strconv" @@ -44,8 +45,8 @@ func TestUpdateConfig_Empty(t *testing.T) { } } -func TestModelserverProxySettingsUsesStableLocalCredential(t *testing.T) { - got := ModelserverProxySettings("http://127.0.0.1:53452/v1") +func TestModelserverProxySettingsUsesConfiguredLocalCredential(t *testing.T) { + got := ModelserverProxySettings("http://127.0.0.1:53452/v1", "random-local-token") if got.Provider != "modelserver" { t.Fatalf("Provider = %q, want modelserver", got.Provider) } @@ -69,9 +70,12 @@ func TestModelserverProxySettingsUsesStableLocalCredential(t *testing.T) { if strings.Contains(s, "env_key") { t.Fatalf("proxy config should not require an environment variable:\n%s", s) } - if !strings.Contains(s, `experimental_bearer_token = "`+LocalProxyAPIKeyValue+`"`) { + if !strings.Contains(s, `experimental_bearer_token = "random-local-token"`) { t.Fatalf("proxy config missing stable bearer token:\n%s", s) } + if strings.Contains(s, "agentserver-local-proxy") { + t.Fatalf("proxy config contains compiled default token:\n%s", s) + } } func TestHasModelserverDirectConfigRequiresExactDirectProvider(t *testing.T) { @@ -86,7 +90,7 @@ func TestHasModelserverDirectConfigRequiresExactDirectProvider(t *testing.T) { `[model_providers.modelserver]`, `name = "modelserver"`, `base_url = "http://127.0.0.1:53452/v1"`, - `experimental_bearer_token = "agentserver-local-proxy"`, + `experimental_bearer_token = "random-local-token"`, `wire_api = "responses"`, ``, }, "\n")) @@ -149,6 +153,73 @@ base_url = "https://old/v1" } } +func TestUpdateConfigPreservesUnknownModelProviderFields(t *testing.T) { + dir := t.TempDir() + path := filepath.Join(dir, "config.toml") + writeCodexTestFile(t, path, strings.Join([]string{ + `model_provider = "modelserver"`, + ``, + `[model_providers.modelserver]`, + `name = "old-name"`, + `base_url = "https://old/v1"`, + `env_key = "OLD_KEY"`, + `wire_api = "chat"`, + `custom_header = "keep"`, + ``, + }, "\n")) + + if err := UpdateConfig(path, ModelserverSettings()); err != nil { + t.Fatal(err) + } + + b, _ := os.ReadFile(path) + s := string(b) + for _, want := range []string{ + `base_url = "https://code.ai.cs.ac.cn/v1"`, + `env_key = "OPENAI_API_KEY"`, + `wire_api = "responses"`, + `custom_header = "keep"`, + } { + if !strings.Contains(s, want) { + t.Fatalf("missing %q in:\n%s", want, s) + } + } +} + +func TestUpdateConfigReturnsBackupError(t *testing.T) { + dir := t.TempDir() + path := filepath.Join(dir, "config.toml") + writeCodexTestFile(t, path, `model_provider = "old"`+"\n") + wantErr := errors.New("backup failed") + prev := writeConfigBackup + writeConfigBackup = func(string, []byte) error { + return wantErr + } + t.Cleanup(func() { + writeConfigBackup = prev + }) + + err := UpdateConfig(path, ModelserverSettings()) + if !errors.Is(err, wantErr) { + t.Fatalf("UpdateConfig error=%v, want %v", err, wantErr) + } +} + +func TestUpdateConfigWritesPrivateFileMode(t *testing.T) { + dir := t.TempDir() + path := filepath.Join(dir, "config.toml") + if err := UpdateConfig(path, ModelserverSettings()); err != nil { + t.Fatal(err) + } + info, err := os.Stat(path) + if err != nil { + t.Fatal(err) + } + if got := info.Mode().Perm(); got != 0o600 { + t.Fatalf("mode=%#o, want 0600", got) + } +} + func TestUpdateConfigPrunesOldBackups(t *testing.T) { dir := t.TempDir() path := filepath.Join(dir, "config.toml") @@ -219,6 +290,37 @@ func TestUpdateMCPServerAddsDriverAndKeepsModelConfig(t *testing.T) { } } +func TestUpdateMCPServerPreservesUnknownFieldsInNamedServer(t *testing.T) { + dir := t.TempDir() + path := filepath.Join(dir, "config.toml") + writeCodexTestFile(t, path, strings.Join([]string{ + `[mcp_servers.driver]`, + `command = "old"`, + `args = ["old"]`, + `custom_tools = ["keep"]`, + ``, + }, "\n")) + + if err := UpdateMCPServer(path, "driver", MCPServer{ + Command: "agentserver", + Args: []string{"serve-driver-mcp"}, + }); err != nil { + t.Fatal(err) + } + + b, _ := os.ReadFile(path) + s := string(b) + for _, want := range []string{ + `command = "agentserver"`, + `args = ["serve-driver-mcp"]`, + `custom_tools = ["keep"]`, + } { + if !strings.Contains(s, want) { + t.Fatalf("missing %q in:\n%s", want, s) + } + } +} + func TestRemoveMCPServerRemovesOnlyNamedServer(t *testing.T) { dir := t.TempDir() path := filepath.Join(dir, "config.toml") diff --git a/internal/headless/driver.go b/internal/headless/driver.go index 08403de..f6a5686 100644 --- a/internal/headless/driver.go +++ b/internal/headless/driver.go @@ -8,6 +8,7 @@ import ( "os" osexec "os/exec" "path/filepath" + "strconv" "strings" "github.com/agentserver/agentserver-pkg/internal/agentserver" @@ -16,18 +17,20 @@ import ( "github.com/agentserver/agentserver-pkg/internal/oauth" "github.com/agentserver/agentserver-pkg/internal/paths" "github.com/agentserver/agentserver-pkg/internal/secrets" + "github.com/agentserver/agentserver-pkg/internal/slave" "github.com/agentserver/agentserver-pkg/internal/state" "github.com/agentserver/agentserver-pkg/internal/terminalauth" ) const ( - defaultAgentserverEndpoint = "https://agent.cs.ac.cn" + defaultAgentserverEndpoint = agentserver.DefaultEndpoint driverProxySecretKey = "agentserver_ws_api_key" driverTunnelSecretKey = "agentserver_tunnel_token" ) type AgentRegistrar interface { RegisterAgent(context.Context, string, string, string) (agentserver.AgentRegistration, error) + DeleteAgent(context.Context, string, string) error Whoami(context.Context, string) (agentserver.Identity, error) } @@ -80,7 +83,11 @@ func ServeDriverMCP(ctx context.Context, opts DriverMCPOptions) error { if err != nil { return err } - sessionConfig := filepath.Join(opts.Paths.InstallRoot, "driver-mcp", fmt.Sprintf("driver-%d.yaml", os.Getpid())) + sessionDir := filepath.Join(opts.Paths.InstallRoot, "driver-mcp") + if err := sweepStaleDriverMCPConfigs(sessionDir, os.Getpid()); err != nil { + return err + } + sessionConfig := filepath.Join(sessionDir, fmt.Sprintf("driver-%d.yaml", os.Getpid())) if err := writeDriverConfig(sessionConfig, st, proxyToken, tunnelToken, workDir); err != nil { return err } @@ -94,6 +101,24 @@ func ServeDriverMCP(ctx context.Context, opts DriverMCPOptions) error { return run(ctx, opts.Package.DriverAgent, []string{"serve-mcp", "--config", sessionConfig}) } +func sweepStaleDriverMCPConfigs(dir string, currentPID int) error { + matches, err := filepath.Glob(filepath.Join(dir, "driver-*.yaml")) + if err != nil { + return err + } + for _, match := range matches { + pidText := strings.TrimSuffix(strings.TrimPrefix(filepath.Base(match), "driver-"), ".yaml") + pid, err := strconv.Atoi(pidText) + if err == nil && pid > 0 && pid != currentPID && slave.ProcessExists(pid) { + continue + } + if err := os.Remove(match); err != nil && !errors.Is(err, os.ErrNotExist) { + return err + } + } + return nil +} + func ensureDriver(ctx context.Context, opts DriverOptions, forceLogin bool) error { if opts.Secrets == nil { return errors.New("install driver: secrets store required") @@ -101,10 +126,16 @@ func ensureDriver(ctx context.Context, opts DriverOptions, forceLogin bool) erro opts = defaultDriverOptions(opts) if forceLogin { - if err := registerDriver(ctx, opts); err != nil { + previous := currentRegisteredDriverAgent(opts.Paths, opts.Secrets) + registered, err := registerDriver(ctx, opts) + if err != nil { + return err + } + if err := refreshDriverFiles(opts.Paths, opts.Package, opts.Secrets); err != nil { return err } - return refreshDriverFiles(opts.Paths, opts.Package, opts.Secrets) + deleteDriverAgentBestEffort(ctx, opts, previous, registered) + return nil } registered, err := driverRegistered(opts.Paths, opts.Secrets) @@ -112,7 +143,7 @@ func ensureDriver(ctx context.Context, opts DriverOptions, forceLogin bool) erro return err } if !registered { - if err := registerDriver(ctx, opts); err != nil { + if _, err := registerDriver(ctx, opts); err != nil { return err } } else if err := repairRegisteredDriverState(ctx, opts); err != nil { @@ -166,43 +197,53 @@ func defaultASBaseURL(cfg oauth.Config) string { return defaultAgentserverEndpoint } -func registerDriver(ctx context.Context, opts DriverOptions) error { +type registeredDriverAgent struct { + proxyToken string + sandboxID string +} + +func registerDriver(ctx context.Context, opts DriverOptions) (registeredDriverAgent, error) { computerName := strings.TrimSpace(opts.ComputerName) if computerName == "" { hostname, err := os.Hostname() if err != nil { - return fmt.Errorf("hostname: %w", err) + return registeredDriverAgent{}, fmt.Errorf("hostname: %w", err) } computerName = hostname } ch, err := opts.RequestDeviceCode(ctx, opts.ASOAuth) if err != nil { - return fmt.Errorf("request agentserver device code: %w", err) + return registeredDriverAgent{}, fmt.Errorf("request agentserver device code: %w", err) } terminalauth.PrintChallenge(opts.Stdout, "Agentserver 登录", ch, opts.QR) tok, err := opts.PollToken(ctx, opts.ASOAuth, ch) if err != nil { - return fmt.Errorf("poll agentserver token: %w", err) + return registeredDriverAgent{}, fmt.Errorf("poll agentserver token: %w", err) } reg, err := opts.AS.RegisterAgent(ctx, tok.AccessToken, computerName+"-星池指挥官", "custom") if err != nil { - return fmt.Errorf("register agentserver agent: %w", err) + return registeredDriverAgent{}, fmt.Errorf("register agentserver agent: %w", err) } if reg.ProxyToken == "" || reg.TunnelToken == "" || reg.SandboxID == "" { - return errors.New("register agentserver agent: incomplete registration") + return registeredDriverAgent{}, errors.New("register agentserver agent: incomplete registration") + } + registered := registeredDriverAgent{proxyToken: reg.ProxyToken, sandboxID: reg.SandboxID} + rollback := func(err error) (registeredDriverAgent, error) { + deleteDriverAgentBestEffort(ctx, opts, registered, registeredDriverAgent{}) + return registeredDriverAgent{}, err } workspace := resolveDriverWorkspace(ctx, opts.AS, reg, tok.AccessToken) if workspace.ID == "" { - return errors.New("resolve agentserver workspace: empty workspace id") + return rollback(errors.New("resolve agentserver workspace: empty workspace id")) } if err := opts.Secrets.Set(driverProxySecretKey, reg.ProxyToken); err != nil { - return err + return rollback(err) } if err := opts.Secrets.Set(driverTunnelSecretKey, reg.TunnelToken); err != nil { - return err + return rollback(err) } - return state.NewStore(opts.Paths.StateFile).Update(func(s *state.State) error { + if err := state.NewStore(opts.Paths.StateFile).Update(func(s *state.State) error { s.Agentserver.BaseURL = defaultASBaseURL(opts.ASOAuth) s.Agentserver.SandboxID = reg.SandboxID s.Agentserver.ShortID = defaultString(reg.ShortID, reg.SandboxID) @@ -210,7 +251,10 @@ func registerDriver(ctx context.Context, opts DriverOptions) error { s.Agentserver.WorkspaceName = workspace.Name s.Agentserver.WorkspaceAPIKeySuffix = lastN(reg.ProxyToken, 4) return nil - }) + }); err != nil { + return rollback(err) + } + return registered, nil } func resolveDriverWorkspace(ctx context.Context, as AgentRegistrar, reg agentserver.AgentRegistration, accessToken string) agentserver.Workspace { @@ -234,6 +278,8 @@ func repairRegisteredDriverState(ctx context.Context, opts DriverOptions) error if opts.AS != nil { if identity, err := opts.AS.Whoami(ctx, proxyToken); err == nil { workspace = identity.Workspace + } else if opts.Stdout != nil { + _, _ = fmt.Fprintf(opts.Stdout, "warning: could not refresh agentserver driver workspace: %v\n", err) } } return state.NewStore(opts.Paths.StateFile).Update(func(s *state.State) error { @@ -269,6 +315,34 @@ func driverRegistered(p paths.Paths, sec secrets.Store) (bool, error) { return true, nil } +func currentRegisteredDriverAgent(p paths.Paths, sec secrets.Store) registeredDriverAgent { + st, err := state.NewStore(p.StateFile).Load() + if err != nil { + return registeredDriverAgent{} + } + sandboxID := strings.TrimSpace(st.Agentserver.SandboxID) + if sandboxID == "" { + return registeredDriverAgent{} + } + proxyToken, err := sec.Get(driverProxySecretKey) + if err != nil || strings.TrimSpace(proxyToken) == "" { + return registeredDriverAgent{} + } + return registeredDriverAgent{proxyToken: proxyToken, sandboxID: sandboxID} +} + +func deleteDriverAgentBestEffort(ctx context.Context, opts DriverOptions, old, replacement registeredDriverAgent) { + if old.sandboxID == "" || old.proxyToken == "" || opts.AS == nil { + return + } + if replacement.sandboxID != "" && old.sandboxID == replacement.sandboxID { + return + } + if err := opts.AS.DeleteAgent(ctx, old.proxyToken, old.sandboxID); err != nil && opts.Stdout != nil { + _, _ = fmt.Fprintf(opts.Stdout, "warning: could not delete previous agentserver sandbox %s: %v\n", old.sandboxID, err) + } +} + func secretPresent(sec secrets.Store, key string) (bool, error) { value, err := sec.Get(key) if errors.Is(err, secrets.ErrNotFound) { diff --git a/internal/headless/driver_test.go b/internal/headless/driver_test.go index d4019d1..03dc2aa 100644 --- a/internal/headless/driver_test.go +++ b/internal/headless/driver_test.go @@ -199,6 +199,52 @@ func TestServeDriverMCPRemovesSessionConfigAfterExit(t *testing.T) { } } +func TestServeDriverMCPSweepsStaleSessionConfigsBeforeStart(t *testing.T) { + temp := t.TempDir() + p := testDriverPaths(temp) + pkg := testDriverPackage(temp) + sec := secrets.New(p.SecretsFile) + if err := sec.Set("agentserver_ws_api_key", "proxy-token"); err != nil { + t.Fatal(err) + } + if err := sec.Set("agentserver_tunnel_token", "tunnel-token"); err != nil { + t.Fatal(err) + } + writeDriverState(t, p, state.AgentserverState{ + BaseURL: "https://agent.test", + SandboxID: "sb-1", + ShortID: "abc123", + WorkspaceID: "ws-1", + }) + sessionDir := filepath.Join(p.InstallRoot, "driver-mcp") + if err := os.MkdirAll(sessionDir, 0o755); err != nil { + t.Fatal(err) + } + stale := filepath.Join(sessionDir, "driver-999999.yaml") + if err := os.WriteFile(stale, []byte("proxy_token: stale\n"), 0o600); err != nil { + t.Fatal(err) + } + + err := ServeDriverMCP(context.Background(), DriverMCPOptions{ + Paths: p, + Package: pkg, + Secrets: sec, + WorkDir: temp, + Exec: func(_ context.Context, _ string, args []string) error { + if _, err := os.Stat(stale); !errors.Is(err, os.ErrNotExist) { + t.Fatalf("stale session config stat=%v, want not exist", err) + } + if _, err := os.Stat(args[2]); err != nil { + t.Fatalf("current session config missing: %v", err) + } + return nil + }, + }) + if err != nil { + t.Fatalf("ServeDriverMCP: %v", err) + } +} + func TestSwitchWorkspaceForcesDeviceLogin(t *testing.T) { t.Run("empty state", func(t *testing.T) { temp := t.TempDir() @@ -270,6 +316,74 @@ func TestSwitchWorkspaceForcesDeviceLogin(t *testing.T) { }) } +func TestSwitchWorkspaceDeletesPreviousSandboxAfterSuccessfulRegister(t *testing.T) { + temp := t.TempDir() + p := testDriverPaths(temp) + sec := secrets.New(p.SecretsFile) + if err := sec.Set("agentserver_ws_api_key", "proxy-old"); err != nil { + t.Fatal(err) + } + if err := sec.Set("agentserver_tunnel_token", "tunnel-old"); err != nil { + t.Fatal(err) + } + writeDriverState(t, p, state.AgentserverState{ + BaseURL: "https://agent.old", + SandboxID: "sb-old", + ShortID: "old123", + WorkspaceID: "ws-old", + }) + fakeAS := fakeRegisteredDriverAS() + + err := SwitchWorkspace(context.Background(), DriverOptions{ + Paths: p, + Package: testDriverPackage(temp), + Secrets: sec, + ComputerName: "host", + AS: fakeAS, + ASOAuth: oauth.Config{Endpoint: "https://agent.test"}, + RequestDeviceCode: func(context.Context, oauth.Config) (oauth.DeviceCodeChallenge, error) { + return driverChallenge(), nil + }, + PollToken: driverPollToken, + Stdout: ioDiscard{}, + QR: driverMarkerQR("fake QR"), + }) + if err != nil { + t.Fatalf("SwitchWorkspace: %v", err) + } + if fakeAS.deletedToken != "proxy-old" || fakeAS.deletedSandboxID != "sb-old" { + t.Fatalf("deleted token=%q sandbox=%q, want old registration", fakeAS.deletedToken, fakeAS.deletedSandboxID) + } +} + +func TestInstallDriverRollsBackRegisteredAgentWhenSecretWriteFails(t *testing.T) { + temp := t.TempDir() + p := testDriverPaths(temp) + sec := &failingSetSecrets{err: errors.New("secret store unavailable")} + fakeAS := fakeRegisteredDriverAS() + + err := InstallDriver(context.Background(), DriverOptions{ + Paths: p, + Package: testDriverPackage(temp), + Secrets: sec, + ComputerName: "host", + AS: fakeAS, + ASOAuth: oauth.Config{Endpoint: "https://agent.test"}, + RequestDeviceCode: func(context.Context, oauth.Config) (oauth.DeviceCodeChallenge, error) { + return driverChallenge(), nil + }, + PollToken: driverPollToken, + Stdout: ioDiscard{}, + QR: driverMarkerQR("fake QR"), + }) + if err == nil || !strings.Contains(err.Error(), "secret store unavailable") { + t.Fatalf("InstallDriver error=%v, want secret store unavailable", err) + } + if fakeAS.deletedToken != "proxy-new" || fakeAS.deletedSandboxID != "sb-new" { + t.Fatalf("deleted token=%q sandbox=%q, want new registration rollback", fakeAS.deletedToken, fakeAS.deletedSandboxID) + } +} + func TestInstallDriverSkipsDeviceFlowWhenAlreadyRegistered(t *testing.T) { temp := t.TempDir() p := testDriverPaths(temp) @@ -467,6 +581,46 @@ func TestInstallDriverRepairsMissingWorkspaceIDWithoutDeviceFlow(t *testing.T) { } } +func TestInstallDriverWarnsWhenRepairWhoamiFails(t *testing.T) { + temp := t.TempDir() + p := testDriverPaths(temp) + sec := secrets.New(p.SecretsFile) + if err := sec.Set("agentserver_ws_api_key", "proxy-token"); err != nil { + t.Fatal(err) + } + if err := sec.Set("agentserver_tunnel_token", "tunnel-token"); err != nil { + t.Fatal(err) + } + writeDriverState(t, p, state.AgentserverState{ + BaseURL: "https://agent.test", + SandboxID: "sb-1", + ShortID: "abc123", + WorkspaceID: "ws-old", + }) + fakeAS := &fakeDriverAgentserver{whoamiErr: errors.New("GET /api/agent/whoami: status 401")} + var out bytes.Buffer + + err := InstallDriver(context.Background(), DriverOptions{ + Paths: p, + Package: testDriverPackage(temp), + Secrets: sec, + AS: fakeAS, + ASOAuth: oauth.Config{Endpoint: "https://agent.test"}, + RequestDeviceCode: func(context.Context, oauth.Config) (oauth.DeviceCodeChallenge, error) { + t.Fatal("RequestDeviceCode called for existing registration") + return oauth.DeviceCodeChallenge{}, nil + }, + PollToken: driverPollToken, + Stdout: &out, + }) + if err != nil { + t.Fatalf("InstallDriver: %v", err) + } + if !strings.Contains(out.String(), "warning:") || !strings.Contains(out.String(), "status 401") { + t.Fatalf("warning output missing whoami failure:\n%s", out.String()) + } +} + func TestInstallDriverRepairsStaleWorkspaceFromCurrentSecrets(t *testing.T) { temp := t.TempDir() p := testDriverPaths(temp) @@ -713,12 +867,14 @@ func TestInstallDriverUsesSandboxIDWhenRegistrationOmitsShortID(t *testing.T) { } type fakeDriverAgentserver struct { - reg agentserver.AgentRegistration - whoami agentserver.Identity - whoamiErr error - whoamiToken string - registerName string - registerType string + reg agentserver.AgentRegistration + whoami agentserver.Identity + whoamiErr error + whoamiToken string + registerName string + registerType string + deletedToken string + deletedSandboxID string } func (f *fakeDriverAgentserver) RegisterAgent(_ context.Context, _ string, name, typ string) (agentserver.AgentRegistration, error) { @@ -735,6 +891,12 @@ func (f *fakeDriverAgentserver) Whoami(_ context.Context, token string) (agentse return f.whoami, nil } +func (f *fakeDriverAgentserver) DeleteAgent(_ context.Context, token, sandboxID string) error { + f.deletedToken = token + f.deletedSandboxID = sandboxID + return nil +} + func fakeRegisteredDriverAS() *fakeDriverAgentserver { return &fakeDriverAgentserver{ reg: agentserver.AgentRegistration{ @@ -841,6 +1003,22 @@ func (s *writeableBrokenReadSecrets) Delete(key string) error { return nil } +type failingSetSecrets struct { + err error +} + +func (s *failingSetSecrets) Get(string) (string, error) { + return "", secrets.ErrNotFound +} + +func (s *failingSetSecrets) Set(string, string) error { + return s.err +} + +func (s *failingSetSecrets) Delete(string) error { + return nil +} + func writeTestTarGz(t *testing.T, path string, files map[string]string) { t.Helper() if err := os.MkdirAll(filepath.Dir(path), 0o755); err != nil { diff --git a/internal/headless/runtime.go b/internal/headless/runtime.go index 9f26427..a0cc741 100644 --- a/internal/headless/runtime.go +++ b/internal/headless/runtime.go @@ -6,7 +6,9 @@ import ( "fmt" "os/exec" "path/filepath" + "regexp" "runtime" + "strconv" "github.com/agentserver/agentserver-pkg/internal/codexruntime" "github.com/agentserver/agentserver-pkg/internal/paths" @@ -28,6 +30,7 @@ type CodexResolveOptions struct { Paths paths.Paths Package Package LookPath func(string) (string, error) + CodexVersion func(context.Context, string) (string, error) EnsureRuntime func(context.Context, string, string, string) (string, error) } @@ -77,7 +80,9 @@ func ResolveCodex(ctx context.Context, opts CodexResolveOptions) (CodexRuntime, lookPath = exec.LookPath } if codexPath, err := lookPath("codex"); err == nil { - return CodexRuntime{Path: codexPath, Source: "path"}, nil + if pathCodexSatisfiesManifest(ctx, codexPath, opts) { + return CodexRuntime{Path: codexPath, Source: "path"}, nil + } } if err := validateManagedCodexOptions(opts); err != nil { return CodexRuntime{}, err @@ -112,6 +117,78 @@ func ResolveCodex(ctx context.Context, opts CodexResolveOptions) (CodexRuntime, return CodexRuntime{Path: codexPath, Source: "managed"}, nil } +func pathCodexSatisfiesManifest(ctx context.Context, codexPath string, opts CodexResolveOptions) bool { + manifestPath, err := opts.Package.codexManifestPath(runtime.GOOS, runtime.GOARCH) + if err != nil { + return true + } + manifest, err := codexruntime.LoadManifest(manifestPath) + if err != nil { + return true + } + want, ok := parseVersionTriple(manifest.PinnedVersion) + if !ok { + return true + } + version := opts.CodexVersion + if version == nil { + version = codexVersion + } + out, err := version(ctx, codexPath) + if err != nil { + return false + } + got, ok := parseVersionTriple(out) + if !ok { + return false + } + return compareVersionTriple(got, want) >= 0 +} + +func codexVersion(ctx context.Context, codexPath string) (string, error) { + out, err := exec.CommandContext(ctx, codexPath, "--version").CombinedOutput() + if err != nil { + return "", err + } + return string(out), nil +} + +var versionTriplePattern = regexp.MustCompile(`\d+\.\d+\.\d+`) +var versionTripleDotPattern = regexp.MustCompile(`\.`) + +func parseVersionTriple(s string) ([3]int, bool) { + var zero [3]int + match := versionTriplePattern.FindString(s) + if match == "" { + return zero, false + } + parts := versionTripleDotPattern.Split(match, 3) + if len(parts) != 3 { + return zero, false + } + var out [3]int + for i, part := range parts { + n, err := strconv.Atoi(part) + if err != nil { + return zero, false + } + out[i] = n + } + return out, true +} + +func compareVersionTriple(a, b [3]int) int { + for i := 0; i < 3; i++ { + switch { + case a[i] < b[i]: + return -1 + case a[i] > b[i]: + return 1 + } + } + return 0 +} + func validateManagedCodexOptions(opts CodexResolveOptions) error { if opts.Paths.CodexExePath == "" { return errors.New("CodexExePath required") diff --git a/internal/headless/runtime_test.go b/internal/headless/runtime_test.go index 85d74df..b58d6f4 100644 --- a/internal/headless/runtime_test.go +++ b/internal/headless/runtime_test.go @@ -3,6 +3,7 @@ package headless import ( "context" "errors" + "os" "os/exec" "path/filepath" "runtime" @@ -69,6 +70,69 @@ func TestResolveCodexPrefersPathCodex(t *testing.T) { } } +func TestResolveCodexFallsBackToManagedRuntimeWhenPathCodexIsOlderThanManifest(t *testing.T) { + skipUnsupportedCodexRuntime(t) + + ctx := context.Background() + temp := t.TempDir() + pathCodex := filepath.Join(temp, "path", exeName("codex")) + pkg := PackagePaths(filepath.Join(temp, "pkg", exeName("agentserver"))) + if err := os.MkdirAll(pkg.PackageDir, 0o755); err != nil { + t.Fatal(err) + } + manifestPath, err := pkg.codexManifestPath(runtime.GOOS, runtime.GOARCH) + if err != nil { + t.Fatal(err) + } + if err := os.WriteFile(manifestPath, []byte(`{ + "package": "@openai/codex", + "platform": "linux-x64", + "pinned_version": "0.139.0-linux-x64", + "strip_prefix": "vendor/x/", + "codex_exe": "bin/codex", + "required_files": ["bin/codex"], + "pinned": { + "integrity": "sha512-test", + "shasum": "test", + "urls": ["https://registry.npmjs.org/@openai/codex/-/codex-0.139.0-linux-x64.tgz"] + } +}`), 0o644); err != nil { + t.Fatal(err) + } + managedCodex := filepath.Join(temp, "managed", exeName("codex")) + ensureCalled := false + + got, err := ResolveCodex(ctx, CodexResolveOptions{ + Paths: paths.Paths{ + CodexExePath: filepath.Join(temp, "bin-root", "bin", exeName("codex")), + CacheDir: filepath.Join(temp, "cache"), + }, + Package: pkg, + LookPath: func(name string) (string, error) { + if name != "codex" { + t.Fatalf("LookPath name=%q, want codex", name) + } + return pathCodex, nil + }, + CodexVersion: func(context.Context, string) (string, error) { + return "codex 0.138.0", nil + }, + EnsureRuntime: func(context.Context, string, string, string) (string, error) { + ensureCalled = true + return managedCodex, nil + }, + }) + if err != nil { + t.Fatal(err) + } + if !ensureCalled { + t.Fatal("EnsureRuntime was not called for old PATH codex") + } + if got.Path != managedCodex || got.Source != "managed" { + t.Fatalf("runtime=%+v, want managed %q", got, managedCodex) + } +} + func TestResolveCodexEnsuresManagedRuntimeWhenPathMissing(t *testing.T) { skipUnsupportedCodexRuntime(t) @@ -390,6 +454,15 @@ func TestLinuxCodexManifestsLoad(t *testing.T) { if m.StripPrefix != tt.stripPrefix { t.Fatalf("StripPrefix=%q, want %q", m.StripPrefix, tt.stripPrefix) } + foundCanonical := false + for _, u := range m.Pinned.URLs { + if strings.Contains(u, "registry.npmjs.org/@openai/codex/") { + foundCanonical = true + } + } + if !foundCanonical { + t.Fatalf("%s missing canonical npm fallback URL: %+v", tt.name, m.Pinned.URLs) + } }) } } diff --git a/internal/headless/slave.go b/internal/headless/slave.go index fbbcdc8..7eb239b 100644 --- a/internal/headless/slave.go +++ b/internal/headless/slave.go @@ -105,6 +105,9 @@ func RunSlave(ctx context.Context, opts SlaveOptions) error { if err != nil { return err } + if sl.PID > 0 && slave.ProcessExists(sl.PID) { + return fmt.Errorf("slave already running for %s (pid %d)", sl.Folder, sl.PID) + } if err := slave.WriteConfig(sl, machine, slave.ConfigInput{CodexBin: opts.CodexBin}); err != nil { return err } diff --git a/internal/headless/slave_test.go b/internal/headless/slave_test.go index c0f9c72..9ba4809 100644 --- a/internal/headless/slave_test.go +++ b/internal/headless/slave_test.go @@ -15,6 +15,7 @@ import ( "time" "github.com/agentserver/agentserver-pkg/internal/paths" + "github.com/agentserver/agentserver-pkg/internal/slave" "github.com/agentserver/agentserver-pkg/internal/terminalauth" ) @@ -227,6 +228,56 @@ func TestRunSlaveReusesExistingEntryWithoutPrompt(t *testing.T) { } } +func TestRunSlaveDoesNotStartWhenExistingSlavePIDIsLive(t *testing.T) { + temp := t.TempDir() + repo := filepath.Join(temp, "repo") + if err := os.Mkdir(repo, 0o755); err != nil { + t.Fatal(err) + } + p := testSlavePaths(temp) + machine, err := slave.NewMachineStore(p.MachineFile).Ensure("host") + if err != nil { + t.Fatal(err) + } + reg := slave.NewRegistry(p.SlavesFile, p.SlavesDir) + existing, created, err := reg.EnsureForFolder(machine, slave.CreateInput{Folder: repo, Name: "worker"}) + if err != nil { + t.Fatalf("EnsureForFolder: %v", err) + } + if !created { + t.Fatal("expected setup to create slave") + } + if _, err := reg.Update(existing.ID, func(sl *slave.Slave) error { + sl.Status = slave.StatusRunning + sl.PID = os.Getpid() + return nil + }); err != nil { + t.Fatalf("Update: %v", err) + } + + runnerCalled := false + err = RunSlave(context.Background(), SlaveOptions{ + Paths: p, + Package: Package{SlaveAgent: filepath.Join(temp, "pkg", "slave-agent")}, + WorkDir: repo, + ComputerName: "host", + NamePrompt: func(string) (string, error) { + t.Fatal("NamePrompt called for existing folder") + return "", nil + }, + Runner: &fakeSlaveProcessRunner{onRun: func() { + runnerCalled = true + }}, + Stdout: ioDiscard{}, + }) + if err == nil || !strings.Contains(err.Error(), "already running") { + t.Fatalf("RunSlave error=%v, want already running", err) + } + if runnerCalled { + t.Fatal("runner was called for existing live slave") + } +} + func TestRunSlavePrintsDuplicateAuthURLOnce(t *testing.T) { temp := t.TempDir() repo := filepath.Join(temp, "repo") diff --git a/internal/launchprep/prepare.go b/internal/launchprep/prepare.go index 4ff9dad..45fefa1 100644 --- a/internal/launchprep/prepare.go +++ b/internal/launchprep/prepare.go @@ -29,7 +29,7 @@ func PrepareVSCode(ctx context.Context, in Input) error { if err := vscode.EnsureTerminalOnlyPanelState(in.Paths.VSCodeUserDataDir); err != nil { return fmt.Errorf("prepare VS Code panel state: %w", err) } - if err := codex.UpdateConfig(in.Paths.CodexConfigFile, codex.ModelserverProxySettings(modelproxy.DefaultBaseURL)); err != nil { + if err := codex.UpdateConfig(in.Paths.CodexConfigFile, codex.ModelserverProxySettings(modelproxy.DefaultBaseURL, codex.LegacyLocalProxyAPIKeyValue)); err != nil { return err } if in.EmbeddedVSIXPath == "" { diff --git a/internal/modelaccess/access.go b/internal/modelaccess/access.go index b764639..7eeea46 100644 --- a/internal/modelaccess/access.go +++ b/internal/modelaccess/access.go @@ -31,6 +31,7 @@ type Result struct { type EnsureOptions struct { CodexConfigPath string Secrets secrets.Store + LocalProxyToken string DeviceConfig oauth.Config AuthCodeConfig oauth.AuthCodeConfig @@ -54,21 +55,18 @@ func Ensure(ctx context.Context, opts EnsureOptions) (Result, error) { opts = defaultEnsureOptions(opts) if opts.Env(tokenrefresh.OpenAIAPIKeyEnv) != "" { - directConfig, err := codex.HasModelserverDirectConfig(opts.CodexConfigPath) - if err != nil { + if err := codex.UpdateConfig(opts.CodexConfigPath, codex.ModelserverSettings()); err != nil { return Result{}, err } - if directConfig { - if err := codex.UpdateConfig(opts.CodexConfigPath, codex.ModelserverSettings()); err != nil { - return Result{}, err - } - return Result{Mode: ModeDirectAPIKey}, nil - } + return Result{Mode: ModeDirectAPIKey}, nil } if opts.Secrets == nil { return Result{}, tokenrefresh.ErrNoSecrets } + if opts.LocalProxyToken == "" { + return Result{}, errors.New("modelaccess: local proxy token required") + } if err := ensureProxyCredentials(ctx, opts); err != nil { return Result{}, err } @@ -77,7 +75,7 @@ func Ensure(ctx context.Context, opts EnsureOptions) (Result, error) { return Result{}, err } } - if err := codex.UpdateConfig(opts.CodexConfigPath, codex.ModelserverProxySettings(modelproxy.DefaultBaseURL)); err != nil { + if err := codex.UpdateConfig(opts.CodexConfigPath, codex.ModelserverProxySettings(modelproxy.DefaultBaseURL, opts.LocalProxyToken)); err != nil { return Result{}, err } return Result{Mode: ModeLocalProxy}, nil @@ -121,7 +119,7 @@ func ensureProxyCredentials(ctx context.Context, opts EnsureOptions) error { } return err } - if accessTokenFresh(opts.Secrets, opts.Now()) { + if !accessTokenNeedsRefresh(opts.Secrets, opts.Now()) { return nil } err = refreshOnce(ctx, opts) @@ -131,7 +129,7 @@ func ensureProxyCredentials(ctx context.Context, opts EnsureOptions) error { return err } -func accessTokenFresh(sec secrets.Store, now time.Time) bool { +func accessTokenNeedsRefresh(sec secrets.Store, now time.Time) bool { raw, err := sec.Get(tokenrefresh.AccessTokenExpiresAtKey) if err != nil || raw == "" { return false @@ -140,7 +138,7 @@ func accessTokenFresh(sec secrets.Store, now time.Time) bool { if err != nil { return false } - return expiresAt.After(now.Add(proxyRefreshBefore)) + return !expiresAt.After(now.Add(proxyRefreshBefore)) } func refreshOnce(ctx context.Context, opts EnsureOptions) error { diff --git a/internal/modelaccess/access_test.go b/internal/modelaccess/access_test.go index 5d9a0a4..60878fe 100644 --- a/internal/modelaccess/access_test.go +++ b/internal/modelaccess/access_test.go @@ -89,9 +89,10 @@ func TestEnsureDirectAPIKeyModeDoesNotRequireSecrets(t *testing.T) { ) } -func TestEnsureOpenAIAPIKeyWithoutMatchingDirectConfigUsesProxy(t *testing.T) { +func TestEnsureOpenAIAPIKeyWithoutMatchingDirectConfigRepairsToDirect(t *testing.T) { tmp := t.TempDir() configPath := filepath.Join(tmp, "codex", "config.toml") + writeProxyConfig(t, configPath, "old-proxy-token") sec := secrets.New(filepath.Join(tmp, "secrets.json")) mustSetSecret(t, sec, tokenrefresh.AccessTokenKey, "existing-access") mustSetSecret(t, sec, tokenrefresh.RefreshTokenKey, "existing-refresh") @@ -120,17 +121,17 @@ func TestEnsureOpenAIAPIKeyWithoutMatchingDirectConfigUsesProxy(t *testing.T) { if err != nil { t.Fatalf("Ensure() error = %v", err) } - if result.Mode != ModeLocalProxy { - t.Fatalf("Mode = %q, want %q", result.Mode, ModeLocalProxy) + if result.Mode != ModeDirectAPIKey { + t.Fatalf("Mode = %q, want %q", result.Mode, ModeDirectAPIKey) } - if !daemonStarted { - t.Fatal("StartDaemon was not called") + if daemonStarted { + t.Fatal("StartDaemon was called in direct API key mode") } assertConfigContains(t, configPath, - `base_url = "`+modelproxy.DefaultBaseURL+`"`, - `experimental_bearer_token = "`+codex.LocalProxyAPIKeyValue+`"`, + `base_url = "https://code.ai.cs.ac.cn/v1"`, + `env_key = "OPENAI_API_KEY"`, ) - assertConfigNotContains(t, configPath, `env_key = "OPENAI_API_KEY"`) + assertConfigNotContains(t, configPath, `experimental_bearer_token`) } func TestEnsureProxyModeRunsDeviceLoginWhenRefreshMissing(t *testing.T) { @@ -150,6 +151,7 @@ func TestEnsureProxyModeRunsDeviceLoginWhenRefreshMissing(t *testing.T) { CodexConfigPath: filepath.Join(tmp, "codex", "config.toml"), Secrets: sec, Env: func(string) string { return "" }, + LocalProxyToken: "random-local-token", RequestDeviceCode: func(context.Context, oauth.Config) (oauth.DeviceCodeChallenge, error) { requestedDeviceCode = true return challenge, nil @@ -199,6 +201,7 @@ func TestEnsureProxyModeUsesExistingRefreshTokenWithoutPrompt(t *testing.T) { CodexConfigPath: filepath.Join(tmp, "codex", "config.toml"), Secrets: sec, Env: func(string) string { return "" }, + LocalProxyToken: "random-local-token", RequestDeviceCode: func(context.Context, oauth.Config) (oauth.DeviceCodeChallenge, error) { t.Fatal("RequestDeviceCode called") return oauth.DeviceCodeChallenge{}, nil @@ -226,6 +229,44 @@ func TestEnsureProxyModeUsesExistingRefreshTokenWithoutPrompt(t *testing.T) { assertSecret(t, sec, tokenrefresh.RefreshTokenKey, "existing-refresh") } +func TestEnsureProxyModeKeepsExistingAccessTokenWhenExpiryMissing(t *testing.T) { + tmp := t.TempDir() + sec := secrets.New(filepath.Join(tmp, "secrets.json")) + mustSetSecret(t, sec, tokenrefresh.AccessTokenKey, "existing-access") + mustSetSecret(t, sec, tokenrefresh.RefreshTokenKey, "existing-refresh") + var daemonStarted bool + + result, err := Ensure(context.Background(), EnsureOptions{ + CodexConfigPath: filepath.Join(tmp, "codex", "config.toml"), + Secrets: sec, + Env: func(string) string { return "" }, + LocalProxyToken: "random-local-token", + Refresh: func(context.Context, oauth.AuthCodeConfig, string) (oauth.Token, error) { + t.Fatal("Refresh called despite existing access token with missing expiry") + return oauth.Token{}, nil + }, + RequestDeviceCode: func(context.Context, oauth.Config) (oauth.DeviceCodeChallenge, error) { + t.Fatal("RequestDeviceCode called despite existing access token") + return oauth.DeviceCodeChallenge{}, nil + }, + StartDaemon: func(context.Context) error { + daemonStarted = true + return nil + }, + Now: fixedNow, + }) + if err != nil { + t.Fatalf("Ensure() error = %v", err) + } + if result.Mode != ModeLocalProxy { + t.Fatalf("Mode = %q, want %q", result.Mode, ModeLocalProxy) + } + if !daemonStarted { + t.Fatal("StartDaemon was not called") + } + assertSecret(t, sec, tokenrefresh.AccessTokenKey, "existing-access") +} + func TestEnsureProxyModeRefreshesExpiredAccessTokenBeforeReturning(t *testing.T) { tmp := t.TempDir() sec := secrets.New(filepath.Join(tmp, "secrets.json")) @@ -238,6 +279,7 @@ func TestEnsureProxyModeRefreshesExpiredAccessTokenBeforeReturning(t *testing.T) CodexConfigPath: filepath.Join(tmp, "codex", "config.toml"), Secrets: sec, Env: func(string) string { return "" }, + LocalProxyToken: "random-local-token", Refresh: func(_ context.Context, _ oauth.AuthCodeConfig, refreshToken string) (oauth.Token, error) { refreshed = true if refreshToken != "existing-refresh" { @@ -279,6 +321,7 @@ func TestEnsureProxyModeRunsDeviceLoginWhenExpiredAccessRefreshNeedsReauth(t *te CodexConfigPath: filepath.Join(tmp, "codex", "config.toml"), Secrets: sec, Env: func(string) string { return "" }, + LocalProxyToken: "random-local-token", Refresh: func(_ context.Context, _ oauth.AuthCodeConfig, refreshToken string) (oauth.Token, error) { refreshed = true return oauth.Token{}, fmt.Errorf("refresh failed: %w", oauth.ErrInvalidGrant) @@ -316,6 +359,7 @@ func TestEnsureProxyModeRefreshesWhenAccessTokenMissing(t *testing.T) { CodexConfigPath: filepath.Join(tmp, "codex", "config.toml"), Secrets: sec, Env: func(string) string { return "" }, + LocalProxyToken: "random-local-token", Refresh: func(_ context.Context, _ oauth.AuthCodeConfig, refreshToken string) (oauth.Token, error) { refreshed = true if refreshToken != "existing-refresh" { @@ -359,6 +403,7 @@ func TestEnsureProxyModeRunsDeviceLoginWhenRefreshNeedsReauth(t *testing.T) { CodexConfigPath: filepath.Join(tmp, "codex", "config.toml"), Secrets: sec, Env: func(string) string { return "" }, + LocalProxyToken: "random-local-token", Refresh: func(_ context.Context, _ oauth.AuthCodeConfig, refreshToken string) (oauth.Token, error) { refreshed = true if refreshToken != "existing-refresh" { @@ -408,6 +453,7 @@ func TestEnsureProxyModePropagatesRefreshErrorWithoutPrompt(t *testing.T) { CodexConfigPath: filepath.Join(tmp, "codex", "config.toml"), Secrets: sec, Env: func(string) string { return "" }, + LocalProxyToken: "random-local-token", Refresh: func(_ context.Context, _ oauth.AuthCodeConfig, refreshToken string) (oauth.Token, error) { refreshed = true if refreshToken != "existing-refresh" { @@ -452,6 +498,7 @@ func TestEnsureRerunsLoginWhenReauthFlagSet(t *testing.T) { CodexConfigPath: filepath.Join(tmp, "codex", "config.toml"), Secrets: sec, Env: func(string) string { return "" }, + LocalProxyToken: "random-local-token", RequestDeviceCode: func(context.Context, oauth.Config) (oauth.DeviceCodeChallenge, error) { requestedDeviceCode = true return oauth.DeviceCodeChallenge{DeviceCode: "device", UserCode: "code", ExpiresIn: 600}, nil @@ -489,6 +536,24 @@ func TestEnsureRequiresSecrets(t *testing.T) { } } +func TestEnsureProxyModeRequiresLocalProxyToken(t *testing.T) { + tmp := t.TempDir() + sec := secrets.New(filepath.Join(tmp, "secrets.json")) + mustSetSecret(t, sec, tokenrefresh.AccessTokenKey, "existing-access") + mustSetSecret(t, sec, tokenrefresh.RefreshTokenKey, "existing-refresh") + mustSetSecret(t, sec, tokenrefresh.AccessTokenExpiresAtKey, fixedNow().Add(time.Hour).Format(time.RFC3339)) + + _, err := Ensure(context.Background(), EnsureOptions{ + CodexConfigPath: filepath.Join(tmp, "codex", "config.toml"), + Secrets: sec, + Env: func(string) string { return "" }, + Now: fixedNow, + }) + if err == nil || !strings.Contains(err.Error(), "local proxy token") { + t.Fatalf("Ensure() error = %v, want local proxy token error", err) + } +} + func TestEnsureProxyModeIgnoresSetEnvFailure(t *testing.T) { tmp := t.TempDir() configPath := filepath.Join(tmp, "codex", "config.toml") @@ -504,6 +569,7 @@ func TestEnsureProxyModeIgnoresSetEnvFailure(t *testing.T) { CodexConfigPath: configPath, Secrets: sec, Env: func(string) string { return "" }, + LocalProxyToken: "random-local-token", SetEnv: func(string, string) error { return setEnvErr }, @@ -544,6 +610,7 @@ func TestEnsureProxyModeIgnoresPersistEnvFailure(t *testing.T) { CodexConfigPath: configPath, Secrets: sec, Env: func(string) string { return "" }, + LocalProxyToken: "random-local-token", SetEnv: func(string, string) error { return nil }, PersistEnv: func(string, string) error { return persistErr @@ -580,6 +647,7 @@ func TestEnsureProxyModeDoesNotWriteProxyConfigWhenStartDaemonFails(t *testing.T CodexConfigPath: configPath, Secrets: sec, Env: func(string) string { return "" }, + LocalProxyToken: "random-local-token", SetEnv: func(string, string) error { return nil }, PersistEnv: func(string, string) error { return nil }, StartDaemon: func(context.Context) error { @@ -606,6 +674,7 @@ func TestProxySettingsWrittenInProxyMode(t *testing.T) { CodexConfigPath: filepath.Join(tmp, "codex", "config.toml"), Secrets: sec, Env: func(string) string { return "" }, + LocalProxyToken: "random-local-token", SetEnv: func(key, value string) error { setEnvCalled = true return nil @@ -628,7 +697,7 @@ func TestProxySettingsWrittenInProxyMode(t *testing.T) { } assertConfigContains(t, filepath.Join(tmp, "codex", "config.toml"), `base_url = "`+modelproxy.DefaultBaseURL+`"`, - `experimental_bearer_token = "`+codex.LocalProxyAPIKeyValue+`"`, + `experimental_bearer_token = "random-local-token"`, ) assertConfigNotContains(t, filepath.Join(tmp, "codex", "config.toml"), `env_key = "`) if setEnvCalled { @@ -646,6 +715,13 @@ func writeDirectConfig(t *testing.T, path string) { } } +func writeProxyConfig(t *testing.T, path, token string) { + t.Helper() + if err := codex.UpdateConfig(path, codex.ModelserverProxySettings(modelproxy.DefaultBaseURL, token)); err != nil { + t.Fatalf("UpdateConfig() error = %v", err) + } +} + func fixedNow() time.Time { return time.Date(2026, 6, 13, 1, 2, 3, 0, time.UTC) } diff --git a/internal/modelaccess/daemon.go b/internal/modelaccess/daemon.go index d1f2e1d..e466ae7 100644 --- a/internal/modelaccess/daemon.go +++ b/internal/modelaccess/daemon.go @@ -33,6 +33,7 @@ var ( type DaemonOptions struct { Secrets secrets.Store OAuth oauth.AuthCodeConfig + LocalProxyToken string ProxyAddr string ProxyUpstreamBaseURL string LockPath string @@ -76,9 +77,10 @@ func RunDaemon(ctx context.Context, opts DaemonOptions) error { go func() { defer wg.Done() resultCh <- daemonResult{name: "proxy", err: modelproxy.ListenAndServe(ctx, modelproxy.ServerOptions{ - Addr: opts.ProxyAddr, - Secrets: opts.Secrets, - UpstreamBaseURL: opts.ProxyUpstreamBaseURL, + Addr: opts.ProxyAddr, + Secrets: opts.Secrets, + UpstreamBaseURL: opts.ProxyUpstreamBaseURL, + LocalBearerToken: opts.LocalProxyToken, })} }() @@ -122,6 +124,7 @@ func RunDaemon(ctx context.Context, opts DaemonOptions) error { type EnsureDaemonOptions struct { ExePath string ProxyBaseURL string + LogPath string HealthCheck func(context.Context, string) bool StartProcess func(*exec.Cmd) error } @@ -148,7 +151,7 @@ func EnsureDaemon(ctx context.Context, opts EnsureDaemonOptions) error { return err } cmd := newDaemonCommand(opts.ExePath, "model-proxy-daemon") - cleanup, err := configureDaemonProcess(cmd) + cleanup, err := configureDaemonProcess(cmd, opts.LogPath) if err != nil { return err } @@ -159,6 +162,9 @@ func EnsureDaemon(ctx context.Context, opts EnsureDaemonOptions) error { } if err := startProcess(cmd); err != nil { if strings.Contains(strings.ToLower(err.Error()), "address already in use") { + if healthCheck(ctx, baseURL) { + return nil + } return errors.Join(ErrProxyUnavailable, err) } return err diff --git a/internal/modelaccess/daemon_process_unix.go b/internal/modelaccess/daemon_process_unix.go index f64b86e..b016b4d 100644 --- a/internal/modelaccess/daemon_process_unix.go +++ b/internal/modelaccess/daemon_process_unix.go @@ -5,17 +5,35 @@ package modelaccess import ( "os" "os/exec" + "path/filepath" "syscall" ) -func configureDaemonProcess(cmd *exec.Cmd) (func(), error) { +func configureDaemonProcess(cmd *exec.Cmd, logPath string) (func(), error) { devNull, err := os.OpenFile(os.DevNull, os.O_RDWR, 0) if err != nil { return nil, err } + logFile := devNull + if logPath != "" { + if err := os.MkdirAll(filepath.Dir(logPath), 0o700); err != nil { + _ = devNull.Close() + return nil, err + } + logFile, err = os.OpenFile(logPath, os.O_CREATE|os.O_WRONLY|os.O_APPEND, 0o600) + if err != nil { + _ = devNull.Close() + return nil, err + } + } cmd.Stdin = devNull - cmd.Stdout = devNull - cmd.Stderr = devNull + cmd.Stdout = logFile + cmd.Stderr = logFile cmd.SysProcAttr = &syscall.SysProcAttr{Setsid: true} - return func() { _ = devNull.Close() }, nil + return func() { + _ = devNull.Close() + if logFile != devNull { + _ = logFile.Close() + } + }, nil } diff --git a/internal/modelaccess/daemon_process_windows.go b/internal/modelaccess/daemon_process_windows.go index 005332a..5a72ede 100644 --- a/internal/modelaccess/daemon_process_windows.go +++ b/internal/modelaccess/daemon_process_windows.go @@ -3,15 +3,35 @@ package modelaccess import ( + "os" "os/exec" + "path/filepath" "github.com/agentserver/agentserver-pkg/internal/process" ) -func configureDaemonProcess(cmd *exec.Cmd) (func(), error) { +func configureDaemonProcess(cmd *exec.Cmd, logPath string) (func(), error) { cmd.Stdin = nil - cmd.Stdout = nil - cmd.Stderr = nil + var logFile *os.File + if logPath != "" { + if err := os.MkdirAll(filepath.Dir(logPath), 0o700); err != nil { + return nil, err + } + var err error + logFile, err = os.OpenFile(logPath, os.O_CREATE|os.O_WRONLY|os.O_APPEND, 0o600) + if err != nil { + return nil, err + } + cmd.Stdout = logFile + cmd.Stderr = logFile + } else { + cmd.Stdout = nil + cmd.Stderr = nil + } process.HideWindow(cmd) - return func() {}, nil + return func() { + if logFile != nil { + _ = logFile.Close() + } + }, nil } diff --git a/internal/modelaccess/daemon_test.go b/internal/modelaccess/daemon_test.go index 525ac6f..a81e9fa 100644 --- a/internal/modelaccess/daemon_test.go +++ b/internal/modelaccess/daemon_test.go @@ -11,7 +11,6 @@ import ( "testing" "time" - "github.com/agentserver/agentserver-pkg/internal/codex" "github.com/agentserver/agentserver-pkg/internal/modelproxy" "github.com/agentserver/agentserver-pkg/internal/oauth" "github.com/agentserver/agentserver-pkg/internal/secrets" @@ -41,6 +40,7 @@ func TestRunDaemonServesLocalModelProxyAndKeepsServingWithoutRefreshToken(t *tes errCh <- RunDaemon(ctx, DaemonOptions{ Secrets: sec, OAuth: oauth.AuthCodeConfig{ClientID: "client-x"}, + LocalProxyToken: "random-local-token", ProxyAddr: addr, ProxyUpstreamBaseURL: upstream.URL + "/v1", LockPath: lockPath, @@ -60,7 +60,7 @@ func TestRunDaemonServesLocalModelProxyAndKeepsServingWithoutRefreshToken(t *tes if err != nil { t.Fatal(err) } - req.Header.Set("Authorization", "Bearer "+codex.LocalProxyAPIKeyValue) + req.Header.Set("Authorization", "Bearer random-local-token") resp, err := http.DefaultClient.Do(req) if err != nil { t.Fatalf("proxy request: %v", err) @@ -118,9 +118,10 @@ func TestRunDaemonStopsRefreshWhenProxyStartupFails(t *testing.T) { ctx, cancel := context.WithCancel(context.Background()) defer cancel() err = RunDaemon(ctx, DaemonOptions{ - Secrets: sec, - OAuth: oauth.AuthCodeConfig{ClientID: "client-x"}, - ProxyAddr: ln.Addr().String(), + Secrets: sec, + OAuth: oauth.AuthCodeConfig{ClientID: "client-x"}, + LocalProxyToken: "random-local-token", + ProxyAddr: ln.Addr().String(), }) if err == nil { t.Fatal("RunDaemon returned nil, want proxy startup error") @@ -161,9 +162,10 @@ func TestRunDaemonStopsProxyWhenRefreshReturnsHardError(t *testing.T) { ctx, cancel := context.WithCancel(context.Background()) defer cancel() err := RunDaemon(ctx, DaemonOptions{ - Secrets: sec, - OAuth: oauth.AuthCodeConfig{ClientID: "client-x"}, - ProxyAddr: addr, + Secrets: sec, + OAuth: oauth.AuthCodeConfig{ClientID: "client-x"}, + LocalProxyToken: "random-local-token", + ProxyAddr: addr, }) if !errors.Is(err, refreshErr) { t.Fatalf("RunDaemon err=%v, want %v", err, refreshErr) @@ -194,9 +196,10 @@ func TestRunDaemonReturnsNilWhenDeadlineExpiresFromRefreshBranch(t *testing.T) { ctx, cancel := context.WithTimeout(context.Background(), time.Nanosecond) time.Sleep(time.Millisecond) err := RunDaemon(ctx, DaemonOptions{ - Secrets: sec, - OAuth: oauth.AuthCodeConfig{ClientID: "client-x"}, - ProxyAddr: freeTCPAddr(t), + Secrets: sec, + OAuth: oauth.AuthCodeConfig{ClientID: "client-x"}, + LocalProxyToken: "random-local-token", + ProxyAddr: freeTCPAddr(t), }) cancel() if err != nil { @@ -314,7 +317,25 @@ func TestEnsureDaemonReusesHealthyExistingProxy(t *testing.T) { } } -func TestEnsureDaemonReportsProxyUnavailableWhenPortOccupied(t *testing.T) { +func TestEnsureDaemonRechecksHealthWhenPortOccupied(t *testing.T) { + healthChecks := 0 + err := EnsureDaemon(context.Background(), EnsureDaemonOptions{ + ExePath: "/opt/agentserver/agentserver", + ProxyBaseURL: "http://127.0.0.1:1", + HealthCheck: func(context.Context, string) bool { + healthChecks++ + return healthChecks > 1 + }, + StartProcess: func(*exec.Cmd) error { + return errors.New("listen tcp 127.0.0.1:53452: bind: address already in use") + }, + }) + if err != nil { + t.Fatalf("EnsureDaemon err=%v, want nil after health recheck", err) + } +} + +func TestEnsureDaemonReportsProxyUnavailableWhenPortOccupiedAndStillUnhealthy(t *testing.T) { err := EnsureDaemon(context.Background(), EnsureDaemonOptions{ ExePath: "/opt/agentserver/agentserver", ProxyBaseURL: "http://127.0.0.1:1", diff --git a/internal/modelaccess/daemon_unix_test.go b/internal/modelaccess/daemon_unix_test.go index 8497fd2..2199895 100644 --- a/internal/modelaccess/daemon_unix_test.go +++ b/internal/modelaccess/daemon_unix_test.go @@ -6,14 +6,17 @@ import ( "context" "os" "os/exec" + "path/filepath" "testing" ) func TestEnsureDaemonDetachesProxyDaemonOnUnix(t *testing.T) { healthy := false + logPath := filepath.Join(t.TempDir(), "logs", "model-proxy-daemon.log") err := EnsureDaemon(context.Background(), EnsureDaemonOptions{ ExePath: "/opt/agentserver/agentserver", ProxyBaseURL: "http://127.0.0.1:1", + LogPath: logPath, HealthCheck: func(context.Context, string) bool { return healthy }, @@ -22,8 +25,8 @@ func TestEnsureDaemonDetachesProxyDaemonOnUnix(t *testing.T) { t.Fatalf("SysProcAttr=%#v, want Setsid", cmd.SysProcAttr) } assertDevNullFile(t, "Stdin", cmd.Stdin) - assertDevNullFile(t, "Stdout", cmd.Stdout) - assertDevNullFile(t, "Stderr", cmd.Stderr) + assertNamedFile(t, "Stdout", cmd.Stdout, logPath) + assertNamedFile(t, "Stderr", cmd.Stderr, logPath) healthy = true return nil }, @@ -34,12 +37,17 @@ func TestEnsureDaemonDetachesProxyDaemonOnUnix(t *testing.T) { } func assertDevNullFile(t *testing.T, name string, got any) { + t.Helper() + assertNamedFile(t, name, got, os.DevNull) +} + +func assertNamedFile(t *testing.T, name string, got any, want string) { t.Helper() f, ok := got.(*os.File) if !ok { t.Fatalf("%s=%T, want *os.File", name, got) } - if f.Name() != os.DevNull { - t.Fatalf("%s file=%q, want %q", name, f.Name(), os.DevNull) + if f.Name() != want { + t.Fatalf("%s file=%q, want %q", name, f.Name(), want) } } diff --git a/internal/modelaccess/proxy_token.go b/internal/modelaccess/proxy_token.go new file mode 100644 index 0000000..7d58e28 --- /dev/null +++ b/internal/modelaccess/proxy_token.go @@ -0,0 +1,81 @@ +package modelaccess + +import ( + "crypto/rand" + "encoding/base64" + "errors" + "fmt" + "os" + "path/filepath" + "strings" +) + +const localProxyTokenBytes = 32 + +func DefaultLocalProxyTokenPath(installRoot string) string { + return filepath.Join(installRoot, "proxy-token") +} + +func EnsureLocalProxyToken(path string) (string, error) { + if token, err := readLocalProxyToken(path); err == nil { + return token, nil + } else if !errors.Is(err, os.ErrNotExist) { + return "", err + } + token, err := generateLocalProxyToken() + if err != nil { + return "", err + } + if err := writeNewLocalProxyToken(path, token); err != nil { + if errors.Is(err, os.ErrExist) { + return readLocalProxyToken(path) + } + return "", err + } + return token, nil +} + +func readLocalProxyToken(path string) (string, error) { + b, err := os.ReadFile(path) + if err != nil { + return "", err + } + token := strings.TrimSpace(string(b)) + if token == "" { + return "", fmt.Errorf("modelaccess: local proxy token file is empty: %s", path) + } + return token, nil +} + +func generateLocalProxyToken() (string, error) { + var b [localProxyTokenBytes]byte + if _, err := rand.Read(b[:]); err != nil { + return "", err + } + return base64.RawURLEncoding.EncodeToString(b[:]), nil +} + +func writeNewLocalProxyToken(path, token string) error { + if err := os.MkdirAll(filepath.Dir(path), 0o700); err != nil { + return err + } + f, err := os.OpenFile(path, os.O_WRONLY|os.O_CREATE|os.O_EXCL, 0o600) + if err != nil { + return err + } + if _, err := f.WriteString(token + "\n"); err != nil { + _ = f.Close() + _ = os.Remove(path) + return err + } + if err := f.Sync(); err != nil { + _ = f.Close() + _ = os.Remove(path) + return err + } + if err := f.Close(); err != nil { + _ = os.Remove(path) + return err + } + return nil +} diff --git a/internal/modelproxy/proxy.go b/internal/modelproxy/proxy.go index fb5276f..f94b1ff 100644 --- a/internal/modelproxy/proxy.go +++ b/internal/modelproxy/proxy.go @@ -13,7 +13,6 @@ import ( "strings" "time" - "github.com/agentserver/agentserver-pkg/internal/codex" "github.com/agentserver/agentserver-pkg/internal/secrets" "github.com/agentserver/agentserver-pkg/internal/tokenrefresh" ) @@ -24,25 +23,48 @@ const ( DefaultUpstreamBaseURL = "https://code.ai.cs.ac.cn/v1" HealthPath = "/agentserver/model-proxy/health" + + MaxRequestBodyBytes = 32 << 20 + + maxHeaderBytes = 64 << 10 + readHeaderTimeout = 10 * time.Second + idleTimeout = 60 * time.Second ) +var hopByHopHeaders = []string{ + "Connection", + "Keep-Alive", + "Proxy-Authenticate", + "Proxy-Authorization", + "TE", + "Trailer", + "Transfer-Encoding", + "Upgrade", +} + type Options struct { - Secrets secrets.Store - UpstreamBaseURL string - Transport http.RoundTripper + Secrets secrets.Store + UpstreamBaseURL string + LocalBearerToken string + Transport http.RoundTripper } type ServerOptions struct { - Addr string - Secrets secrets.Store - UpstreamBaseURL string - Transport http.RoundTripper + Addr string + Secrets secrets.Store + UpstreamBaseURL string + LocalBearerToken string + Transport http.RoundTripper } func NewHandler(opts Options) (http.Handler, error) { if opts.Secrets == nil { return nil, errors.New("modelproxy: secrets store required") } + localBearerToken := strings.TrimSpace(opts.LocalBearerToken) + if localBearerToken == "" { + return nil, errors.New("modelproxy: local bearer token required") + } upstreamRaw := opts.UpstreamBaseURL if upstreamRaw == "" { upstreamRaw = DefaultUpstreamBaseURL @@ -71,10 +93,14 @@ func NewHandler(opts Options) (http.Handler, error) { w.WriteHeader(http.StatusNoContent) return } - if !validLocalBearer(r.Header.Get("Authorization")) { + if !validLocalBearer(r.Header.Get("Authorization"), localBearerToken) { http.Error(w, "local model proxy authorization required", http.StatusUnauthorized) return } + if r.ContentLength > MaxRequestBodyBytes { + http.Error(w, "request body too large", http.StatusRequestEntityTooLarge) + return + } token, err := opts.Secrets.Get(tokenrefresh.AccessTokenKey) if err != nil || token == "" { http.Error(w, "modelserver login required", http.StatusUnauthorized) @@ -82,17 +108,32 @@ func NewHandler(opts Options) (http.Handler, error) { } r2 := r.Clone(r.Context()) r2.Header = r.Header.Clone() + stripHopByHopHeaders(r2.Header) r2.Header.Set("Authorization", "Bearer "+token) + r2.Body = http.MaxBytesReader(w, r.Body, MaxRequestBodyBytes) proxy.ServeHTTP(w, r2) }), nil } -func validLocalBearer(auth string) bool { +func stripHopByHopHeaders(h http.Header) { + for _, value := range h.Values("Connection") { + for _, name := range strings.Split(value, ",") { + if name = strings.TrimSpace(name); name != "" { + h.Del(name) + } + } + } + for _, name := range hopByHopHeaders { + h.Del(name) + } +} + +func validLocalBearer(auth, token string) bool { parts := strings.Fields(auth) if len(parts) != 2 || !strings.EqualFold(parts[0], "Bearer") { return false } - return subtle.ConstantTimeCompare([]byte(parts[1]), []byte(codex.LocalProxyAPIKeyValue)) == 1 + return subtle.ConstantTimeCompare([]byte(parts[1]), []byte(token)) == 1 } func ListenAndServe(ctx context.Context, opts ServerOptions) error { @@ -104,9 +145,10 @@ func ListenAndServe(ctx context.Context, opts ServerOptions) error { addr = DefaultListenAddr } handler, err := NewHandler(Options{ - Secrets: opts.Secrets, - UpstreamBaseURL: opts.UpstreamBaseURL, - Transport: opts.Transport, + Secrets: opts.Secrets, + UpstreamBaseURL: opts.UpstreamBaseURL, + LocalBearerToken: opts.LocalBearerToken, + Transport: opts.Transport, }) if err != nil { return err @@ -115,7 +157,12 @@ func ListenAndServe(ctx context.Context, opts ServerOptions) error { if err != nil { return err } - srv := &http.Server{Handler: handler} + srv := &http.Server{ + Handler: handler, + ReadHeaderTimeout: readHeaderTimeout, + IdleTimeout: idleTimeout, + MaxHeaderBytes: maxHeaderBytes, + } go func() { <-ctx.Done() shutdownCtx, cancel := context.WithTimeout(context.Background(), 5*time.Second) diff --git a/internal/modelproxy/proxy_test.go b/internal/modelproxy/proxy_test.go index 8bbd1a8..ccf3a6a 100644 --- a/internal/modelproxy/proxy_test.go +++ b/internal/modelproxy/proxy_test.go @@ -1,12 +1,12 @@ package modelproxy import ( + "bytes" "io" "net/http" "net/http/httptest" "testing" - "github.com/agentserver/agentserver-pkg/internal/codex" "github.com/agentserver/agentserver-pkg/internal/secrets" "github.com/agentserver/agentserver-pkg/internal/tokenrefresh" ) @@ -55,8 +55,9 @@ func TestProxyLoadsLatestAccessTokenForEveryRequest(t *testing.T) { defer upstream.Close() handler, err := NewHandler(Options{ - Secrets: sec, - UpstreamBaseURL: upstream.URL + "/v1", + Secrets: sec, + UpstreamBaseURL: upstream.URL + "/v1", + LocalBearerToken: "random-local-token", }) if err != nil { t.Fatalf("NewHandler: %v", err) @@ -68,7 +69,7 @@ func TestProxyLoadsLatestAccessTokenForEveryRequest(t *testing.T) { if err != nil { t.Fatal(err) } - req.Header.Set("Authorization", "Bearer "+codex.LocalProxyAPIKeyValue) + req.Header.Set("Authorization", "Bearer random-local-token") req.Header.Set("Content-Type", "text/plain") resp, err := http.DefaultClient.Do(req) if err != nil { @@ -89,7 +90,7 @@ func TestProxyLoadsLatestAccessTokenForEveryRequest(t *testing.T) { if err != nil { t.Fatal(err) } - req.Header.Set("Authorization", "Bearer "+codex.LocalProxyAPIKeyValue) + req.Header.Set("Authorization", "Bearer random-local-token") resp, err = http.DefaultClient.Do(req) if err != nil { t.Fatalf("second request: %v", err) @@ -125,8 +126,9 @@ func TestProxyRequiresLocalBearerToken(t *testing.T) { defer upstream.Close() handler, err := NewHandler(Options{ - Secrets: sec, - UpstreamBaseURL: upstream.URL + "/v1", + Secrets: sec, + UpstreamBaseURL: upstream.URL + "/v1", + LocalBearerToken: "random-local-token", }) if err != nil { t.Fatalf("NewHandler: %v", err) @@ -138,7 +140,7 @@ func TestProxyRequiresLocalBearerToken(t *testing.T) { }{ {name: "missing"}, {name: "wrong token", auth: "Bearer wrong"}, - {name: "wrong scheme", auth: "Basic " + codex.LocalProxyAPIKeyValue}, + {name: "wrong scheme", auth: "Basic random-local-token"}, } { t.Run(tt.name, func(t *testing.T) { upstreamCalled = false @@ -161,8 +163,9 @@ func TestProxyRequiresLocalBearerToken(t *testing.T) { func TestProxyHealthDoesNotRequireLocalBearerToken(t *testing.T) { handler, err := NewHandler(Options{ - Secrets: newTestSecrets(), - UpstreamBaseURL: "https://upstream.test/v1", + Secrets: newTestSecrets(), + UpstreamBaseURL: "https://upstream.test/v1", + LocalBearerToken: "random-local-token", }) if err != nil { t.Fatalf("NewHandler: %v", err) @@ -177,6 +180,83 @@ func TestProxyHealthDoesNotRequireLocalBearerToken(t *testing.T) { } } +func TestProxyRejectsOversizedRequestBody(t *testing.T) { + upstreamCalled := false + upstream := httptest.NewServer(http.HandlerFunc(func(w http.ResponseWriter, r *http.Request) { + upstreamCalled = true + })) + defer upstream.Close() + + sec := newTestSecrets() + if err := sec.Set(tokenrefresh.AccessTokenKey, "access"); err != nil { + t.Fatal(err) + } + handler, err := NewHandler(Options{ + Secrets: sec, + UpstreamBaseURL: upstream.URL + "/v1", + LocalBearerToken: "random-local-token", + }) + if err != nil { + t.Fatalf("NewHandler: %v", err) + } + + req := httptest.NewRequest(http.MethodPost, "/v1/responses", bytes.NewReader(make([]byte, MaxRequestBodyBytes+1))) + req.Header.Set("Authorization", "Bearer random-local-token") + rec := httptest.NewRecorder() + handler.ServeHTTP(rec, req) + + if rec.Code != http.StatusRequestEntityTooLarge { + t.Fatalf("status=%d, want %d", rec.Code, http.StatusRequestEntityTooLarge) + } + if upstreamCalled { + t.Fatal("upstream should not receive oversized request") + } +} + +func TestProxyStripsHopByHopAndProxyAuthorizationHeaders(t *testing.T) { + sec := newTestSecrets() + if err := sec.Set(tokenrefresh.AccessTokenKey, "access"); err != nil { + t.Fatal(err) + } + + var got http.Header + upstream := httptest.NewServer(http.HandlerFunc(func(w http.ResponseWriter, r *http.Request) { + got = r.Header.Clone() + w.WriteHeader(http.StatusOK) + })) + defer upstream.Close() + + handler, err := NewHandler(Options{ + Secrets: sec, + UpstreamBaseURL: upstream.URL + "/v1", + LocalBearerToken: "random-local-token", + }) + if err != nil { + t.Fatalf("NewHandler: %v", err) + } + + req := httptest.NewRequest(http.MethodGet, "/v1/models", nil) + req.Header.Set("Authorization", "Bearer random-local-token") + req.Header.Set("Connection", "keep-alive, X-Drop-Me") + req.Header.Set("X-Drop-Me", "drop") + req.Header.Set("Proxy-Authorization", "Basic secret") + req.Header.Set("TE", "trailers") + rec := httptest.NewRecorder() + handler.ServeHTTP(rec, req) + + if rec.Code != http.StatusOK { + t.Fatalf("status=%d, want %d", rec.Code, http.StatusOK) + } + for _, name := range []string{"Connection", "X-Drop-Me", "Proxy-Authorization", "TE"} { + if got.Get(name) != "" { + t.Fatalf("upstream header %s=%q, want stripped", name, got.Get(name)) + } + } + if got.Get("Authorization") != "Bearer access" { + t.Fatalf("upstream Authorization=%q, want modelserver access token", got.Get("Authorization")) + } +} + func TestProxyReturnsUnauthorizedWhenAccessTokenMissing(t *testing.T) { upstreamCalled := false upstream := httptest.NewServer(http.HandlerFunc(func(w http.ResponseWriter, r *http.Request) { @@ -185,15 +265,16 @@ func TestProxyReturnsUnauthorizedWhenAccessTokenMissing(t *testing.T) { defer upstream.Close() handler, err := NewHandler(Options{ - Secrets: newTestSecrets(), - UpstreamBaseURL: upstream.URL + "/v1", + Secrets: newTestSecrets(), + UpstreamBaseURL: upstream.URL + "/v1", + LocalBearerToken: "random-local-token", }) if err != nil { t.Fatalf("NewHandler: %v", err) } req := httptest.NewRequest(http.MethodGet, "/v1/models", nil) - req.Header.Set("Authorization", "Bearer "+codex.LocalProxyAPIKeyValue) + req.Header.Set("Authorization", "Bearer random-local-token") rec := httptest.NewRecorder() handler.ServeHTTP(rec, req) @@ -204,3 +285,13 @@ func TestProxyReturnsUnauthorizedWhenAccessTokenMissing(t *testing.T) { t.Fatal("upstream should not be called without an access token") } } + +func TestProxyRequiresConfiguredLocalBearerToken(t *testing.T) { + _, err := NewHandler(Options{ + Secrets: newTestSecrets(), + UpstreamBaseURL: "https://upstream.test/v1", + }) + if err == nil { + t.Fatal("NewHandler returned nil error without LocalBearerToken") + } +} diff --git a/internal/oauth/devicecode.go b/internal/oauth/devicecode.go index 9463bfd..db0aba3 100644 --- a/internal/oauth/devicecode.go +++ b/internal/oauth/devicecode.go @@ -5,12 +5,17 @@ import ( "encoding/json" "errors" "fmt" + "io" "net/http" "net/url" "strings" "time" ) +const maxDeviceResponseBytes = 1 << 20 + +var deviceHTTPClient = &http.Client{Timeout: 60 * time.Second} + // RequestDeviceCode POSTs to AuthURL and returns the challenge. func RequestDeviceCode(ctx context.Context, cfg Config) (DeviceCodeChallenge, error) { form := url.Values{} @@ -25,7 +30,7 @@ func RequestDeviceCode(ctx context.Context, cfg Config) (DeviceCodeChallenge, er } req.Header.Set("Content-Type", "application/x-www-form-urlencoded") req.Header.Set("Accept", "application/json") - resp, err := http.DefaultClient.Do(req) + resp, err := deviceHTTPClient.Do(req) if err != nil { return DeviceCodeChallenge{}, fmt.Errorf("device auth: %w", err) } @@ -34,7 +39,7 @@ func RequestDeviceCode(ctx context.Context, cfg Config) (DeviceCodeChallenge, er return DeviceCodeChallenge{}, fmt.Errorf("device auth: status %d", resp.StatusCode) } var ch DeviceCodeChallenge - if err := json.NewDecoder(resp.Body).Decode(&ch); err != nil { + if err := json.NewDecoder(io.LimitReader(resp.Body, maxDeviceResponseBytes)).Decode(&ch); err != nil { return DeviceCodeChallenge{}, fmt.Errorf("decode device auth: %w", err) } ch.RetrievedAt = time.Now() @@ -111,21 +116,21 @@ func tokenOnce(ctx context.Context, cfg Config, deviceCode string) (Token, strin } req.Header.Set("Content-Type", "application/x-www-form-urlencoded") req.Header.Set("Accept", "application/json") - resp, err := http.DefaultClient.Do(req) + resp, err := deviceHTTPClient.Do(req) if err != nil { return Token{}, "", fmt.Errorf("token poll: %w", err) } defer resp.Body.Close() if resp.StatusCode == http.StatusOK { var tok Token - if err := json.NewDecoder(resp.Body).Decode(&tok); err != nil { + if err := json.NewDecoder(io.LimitReader(resp.Body, maxDeviceResponseBytes)).Decode(&tok); err != nil { return Token{}, "", err } return tok, "", nil } // 400 with JSON {error: "..."} is normal device-flow signalling var te tokenErr - if err := json.NewDecoder(resp.Body).Decode(&te); err != nil { + if err := json.NewDecoder(io.LimitReader(resp.Body, maxDeviceResponseBytes)).Decode(&te); err != nil { return Token{}, "", fmt.Errorf("token poll: status %d", resp.StatusCode) } if te.Code == "" { diff --git a/internal/oauth/devicecode_test.go b/internal/oauth/devicecode_test.go index d8a942d..ebb1de5 100644 --- a/internal/oauth/devicecode_test.go +++ b/internal/oauth/devicecode_test.go @@ -5,6 +5,7 @@ import ( "encoding/json" "net/http" "net/http/httptest" + "strings" "sync/atomic" "testing" "time" @@ -59,6 +60,28 @@ func TestRequestDeviceCode(t *testing.T) { } } +func TestRequestDeviceCodeRejectsOversizedResponse(t *testing.T) { + srv := httptest.NewServer(http.HandlerFunc(func(w http.ResponseWriter, r *http.Request) { + _, _ = w.Write([]byte(`{"device_code":"` + strings.Repeat("x", maxDeviceResponseBytes) + `"}`)) + })) + defer srv.Close() + + cfg := Config{Endpoint: srv.URL, AuthPath: "/", ClientID: "test-client"} + _, err := RequestDeviceCode(context.Background(), cfg) + if err == nil { + t.Fatal("RequestDeviceCode returned nil error for oversized response") + } +} + +func TestDeviceHTTPClientHasTimeout(t *testing.T) { + if deviceHTTPClient == http.DefaultClient { + t.Fatal("deviceHTTPClient must not use http.DefaultClient") + } + if deviceHTTPClient.Timeout <= 0 { + t.Fatalf("deviceHTTPClient timeout=%v, want positive timeout", deviceHTTPClient.Timeout) + } +} + func TestPollTokenSuccess(t *testing.T) { srv := fakeHydra(t, 2) // succeed on 3rd poll defer srv.Close() diff --git a/internal/slave/file_lock_unix.go b/internal/slave/file_lock_unix.go new file mode 100644 index 0000000..c937292 --- /dev/null +++ b/internal/slave/file_lock_unix.go @@ -0,0 +1,43 @@ +//go:build !windows + +package slave + +import ( + "fmt" + "os" + "path/filepath" + + "golang.org/x/sys/unix" +) + +type registryFileLock struct { + file *os.File +} + +func acquireRegistryFileLock(path string) (*registryFileLock, error) { + if err := os.MkdirAll(filepath.Dir(path), 0o700); err != nil { + return nil, fmt.Errorf("mkdir slave registry lock dir: %w", err) + } + f, err := os.OpenFile(path, os.O_CREATE|os.O_RDWR, 0o600) + if err != nil { + return nil, fmt.Errorf("open slave registry lock: %w", err) + } + if err := unix.Flock(int(f.Fd()), unix.LOCK_EX); err != nil { + _ = f.Close() + return nil, fmt.Errorf("lock slave registry: %w", err) + } + return ®istryFileLock{file: f}, nil +} + +func (l *registryFileLock) close() error { + if l == nil || l.file == nil { + return nil + } + err := unix.Flock(int(l.file.Fd()), unix.LOCK_UN) + closeErr := l.file.Close() + l.file = nil + if err != nil { + return err + } + return closeErr +} diff --git a/internal/slave/file_lock_windows.go b/internal/slave/file_lock_windows.go new file mode 100644 index 0000000..933d8f5 --- /dev/null +++ b/internal/slave/file_lock_windows.go @@ -0,0 +1,45 @@ +//go:build windows + +package slave + +import ( + "fmt" + "os" + "path/filepath" + + "golang.org/x/sys/windows" +) + +type registryFileLock struct { + file *os.File +} + +func acquireRegistryFileLock(path string) (*registryFileLock, error) { + if err := os.MkdirAll(filepath.Dir(path), 0o700); err != nil { + return nil, fmt.Errorf("mkdir slave registry lock dir: %w", err) + } + f, err := os.OpenFile(path, os.O_CREATE|os.O_RDWR, 0o600) + if err != nil { + return nil, fmt.Errorf("open slave registry lock: %w", err) + } + var overlapped windows.Overlapped + if err := windows.LockFileEx(windows.Handle(f.Fd()), windows.LOCKFILE_EXCLUSIVE_LOCK, 0, 1, 0, &overlapped); err != nil { + _ = f.Close() + return nil, fmt.Errorf("lock slave registry: %w", err) + } + return ®istryFileLock{file: f}, nil +} + +func (l *registryFileLock) close() error { + if l == nil || l.file == nil { + return nil + } + var overlapped windows.Overlapped + err := windows.UnlockFileEx(windows.Handle(l.file.Fd()), 0, 1, 0, &overlapped) + closeErr := l.file.Close() + l.file = nil + if err != nil { + return err + } + return closeErr +} diff --git a/internal/slave/process_liveness.go b/internal/slave/process_liveness.go new file mode 100644 index 0000000..f08f15b --- /dev/null +++ b/internal/slave/process_liveness.go @@ -0,0 +1,6 @@ +package slave + +// ProcessExists reports whether pid currently names a live OS process. +func ProcessExists(pid int) bool { + return osProcessExists(pid) +} diff --git a/internal/slave/registry.go b/internal/slave/registry.go index 126894d..02a0328 100644 --- a/internal/slave/registry.go +++ b/internal/slave/registry.go @@ -112,6 +112,10 @@ func NewRegistry(path, slavesDir string) *Registry { return &Registry{path: path, slavesDir: slavesDir} } +func (r *Registry) lockPath() string { + return r.path + ".lock" +} + func (r *Registry) storageDir(id string) (string, error) { if r == nil { return "", fmt.Errorf("slave registry required") @@ -144,6 +148,12 @@ func (r *Registry) List() ([]Slave, error) { r.mu.Lock() defer r.mu.Unlock() + lock, err := acquireRegistryFileLock(r.lockPath()) + if err != nil { + return nil, err + } + defer lock.close() + all, err := r.loadLocked() if err != nil { return nil, err @@ -155,6 +165,12 @@ func (r *Registry) Get(id string) (Slave, error) { r.mu.Lock() defer r.mu.Unlock() + lock, err := acquireRegistryFileLock(r.lockPath()) + if err != nil { + return Slave{}, err + } + defer lock.close() + all, err := r.loadLocked() if err != nil { return Slave{}, err @@ -189,6 +205,12 @@ func (r *Registry) Create(machine Machine, in CreateInput) (Slave, error) { r.mu.Lock() defer r.mu.Unlock() + lock, err := acquireRegistryFileLock(r.lockPath()) + if err != nil { + return Slave{}, err + } + defer lock.close() + all, err := r.loadLocked() if err != nil { return Slave{}, err @@ -238,25 +260,21 @@ func (r *Registry) EnsureForFolder(machine Machine, in CreateInput) (Slave, bool } displayName := machine.ComputerName + "-" + name - all, err := r.List() + r.mu.Lock() + defer r.mu.Unlock() + + lock, err := acquireRegistryFileLock(r.lockPath()) if err != nil { return Slave{}, false, err } - for _, existing := range all { - if storedFolderKey(existing.Folder) == folder { - return existing, false, nil - } - } - - r.mu.Lock() - defer r.mu.Unlock() + defer lock.close() latest, err := r.loadLocked() if err != nil { return Slave{}, false, err } for _, existing := range latest { - if filepath.Clean(existing.Folder) == folder { + if storedFolderKey(existing.Folder) == folder { return existing, false, nil } } @@ -276,6 +294,12 @@ func (r *Registry) Update(id string, fn func(*Slave) error) (Slave, error) { r.mu.Lock() defer r.mu.Unlock() + lock, err := acquireRegistryFileLock(r.lockPath()) + if err != nil { + return Slave{}, err + } + defer lock.close() + all, err := r.loadLocked() if err != nil { return Slave{}, err @@ -300,6 +324,12 @@ func (r *Registry) Delete(id string) (Slave, error) { r.mu.Lock() defer r.mu.Unlock() + lock, err := acquireRegistryFileLock(r.lockPath()) + if err != nil { + return Slave{}, err + } + defer lock.close() + all, err := r.loadLocked() if err != nil { return Slave{}, err @@ -349,8 +379,8 @@ func validateSlaveName(name string) error { if len([]rune(name)) > 20 { return fmt.Errorf("slave name must be at most 20 characters") } - if strings.ContainsAny(name, `\\/:*?"<>|`) { - return fmt.Errorf("slave name contains invalid path characters") + if strings.ContainsAny(name, `\/`) || strings.ContainsRune(name, 0) { + return fmt.Errorf("slave name contains path separators") } return nil } @@ -371,10 +401,10 @@ func (r *Registry) loadLocked() ([]Slave, error) { } func (r *Registry) saveLocked(all []Slave) error { - if err := osMkdirAll(filepath.Dir(r.path), 0o755); err != nil { + if err := osMkdirAll(filepath.Dir(r.path), 0o700); err != nil { return fmt.Errorf("mkdir slave registry dir: %w", err) } - if err := osMkdirAll(r.slavesDir, 0o755); err != nil { + if err := osMkdirAll(r.slavesDir, 0o700); err != nil { return fmt.Errorf("mkdir slaves dir: %w", err) } b, err := json.MarshalIndent(all, "", " ") diff --git a/internal/slave/registry_test.go b/internal/slave/registry_test.go index a17d987..30977c9 100644 --- a/internal/slave/registry_test.go +++ b/internal/slave/registry_test.go @@ -64,6 +64,24 @@ func TestRegistryCreatesSlaveWithCustomImmutableName(t *testing.T) { } } +func TestRegistryAllowsLinuxLegalNameCharacters(t *testing.T) { + dir := t.TempDir() + folder := filepath.Join(dir, "repo") + if err := mkdir(folder); err != nil { + t.Fatal(err) + } + reg := NewRegistry(filepath.Join(dir, "slaves.json"), filepath.Join(dir, "slaves")) + m := Machine{MachineID: "machine-1", ComputerName: "61414-PC"} + + got, err := reg.Create(m, CreateInput{Folder: folder, Name: "gpu:job*?"}) + if err != nil { + t.Fatalf("Create: %v", err) + } + if got.Name != "gpu:job*?" { + t.Fatalf("Name=%q", got.Name) + } +} + func TestRegistryRejectsInvalidCreateInput(t *testing.T) { dir := t.TempDir() reg := NewRegistry(filepath.Join(dir, "slaves.json"), filepath.Join(dir, "slaves")) @@ -427,6 +445,80 @@ func TestRegistryConcurrentEnsureForFolderCreatesOneSlave(t *testing.T) { } } +func TestRegistryConcurrentEnsureForFolderAcrossInstancesCreatesOneSlave(t *testing.T) { + prevProcs := runtime.GOMAXPROCS(0) + if prevProcs < 8 { + runtime.GOMAXPROCS(8) + t.Cleanup(func() { + runtime.GOMAXPROCS(prevProcs) + }) + } + + const attempts = 20 + const workers = 64 + + for attempt := 0; attempt < attempts; attempt++ { + dir := t.TempDir() + folder := filepath.Join(dir, "repo") + if err := mkdir(folder); err != nil { + t.Fatal(err) + } + registryPath := filepath.Join(dir, "slaves.json") + slavesDir := filepath.Join(dir, "slaves") + m := Machine{MachineID: "machine-1", ComputerName: "host"} + start := make(chan struct{}) + results := make([]Slave, workers) + created := make([]bool, workers) + errs := make([]error, workers) + + var wg sync.WaitGroup + wg.Add(workers) + for i := 0; i < workers; i++ { + i := i + go func() { + defer wg.Done() + <-start + reg := NewRegistry(registryPath, slavesDir) + results[i], created[i], errs[i] = reg.EnsureForFolder(m, CreateInput{Folder: folder}) + }() + } + + close(start) + wg.Wait() + + var id string + createCount := 0 + for i, err := range errs { + if err != nil { + t.Fatalf("attempt %d worker %d EnsureForFolder: %v", attempt, i, err) + } + if results[i].ID == "" { + t.Fatalf("attempt %d worker %d empty ID: %+v", attempt, i, results[i]) + } + if id == "" { + id = results[i].ID + } + if results[i].ID != id { + t.Fatalf("attempt %d worker %d ID=%q want %q", attempt, i, results[i].ID, id) + } + if created[i] { + createCount++ + } + } + if createCount != 1 { + t.Fatalf("attempt %d created count=%d want 1", attempt, createCount) + } + + all, err := NewRegistry(registryPath, slavesDir).List() + if err != nil { + t.Fatalf("attempt %d List: %v", attempt, err) + } + if len(all) != 1 { + t.Fatalf("attempt %d registry has %d slaves, want 1: %+v", attempt, len(all), all) + } + } +} + func TestRegistryConcurrentCreatesPreserveAllSlaves(t *testing.T) { prevProcs := runtime.GOMAXPROCS(0) if prevProcs < 8 { @@ -507,6 +599,64 @@ func TestRegistryConcurrentCreatesPreserveAllSlaves(t *testing.T) { } } +func TestRegistryConcurrentCreatesAcrossInstancesPreserveAllSlaves(t *testing.T) { + prevProcs := runtime.GOMAXPROCS(0) + if prevProcs < 8 { + runtime.GOMAXPROCS(8) + t.Cleanup(func() { + runtime.GOMAXPROCS(prevProcs) + }) + } + + const attempts = 20 + const workers = 64 + + for attempt := 0; attempt < attempts; attempt++ { + dir := t.TempDir() + registryPath := filepath.Join(dir, "slaves.json") + slavesDir := filepath.Join(dir, "slaves") + m := Machine{MachineID: "machine-1", ComputerName: "61414-PC"} + start := make(chan struct{}) + errs := make([]error, workers) + + var wg sync.WaitGroup + wg.Add(workers) + for i := 0; i < workers; i++ { + i := i + go func() { + defer wg.Done() + folder := filepath.Join(dir, fmt.Sprintf("repo-%02d", i)) + if err := mkdir(folder); err != nil { + errs[i] = err + return + } + <-start + reg := NewRegistry(registryPath, slavesDir) + _, errs[i] = reg.Create(m, CreateInput{ + Folder: folder, + Name: fmt.Sprintf("worker-%02d", i), + }) + }() + } + + close(start) + wg.Wait() + + for i, err := range errs { + if err != nil { + t.Fatalf("attempt %d worker %d Create: %v", attempt, i, err) + } + } + got, err := NewRegistry(registryPath, slavesDir).List() + if err != nil { + t.Fatalf("attempt %d List: %v", attempt, err) + } + if len(got) != workers { + t.Fatalf("attempt %d saved %d slaves, want %d: %+v", attempt, len(got), workers, got) + } + } +} + func TestRegistrySaveNarrowsExistingRegistryFileMode(t *testing.T) { if runtime.GOOS == "windows" { t.Skip("windows does not preserve POSIX file mode bits") @@ -540,6 +690,35 @@ func TestRegistrySaveNarrowsExistingRegistryFileMode(t *testing.T) { } } +func TestRegistrySaveCreatesPrivateStateDirs(t *testing.T) { + if runtime.GOOS == "windows" { + t.Skip("windows does not preserve POSIX file mode bits") + } + + dir := t.TempDir() + folder := filepath.Join(dir, "repo") + if err := mkdir(folder); err != nil { + t.Fatal(err) + } + stateDir := filepath.Join(dir, "state") + slavesDir := filepath.Join(stateDir, "slaves") + reg := NewRegistry(filepath.Join(stateDir, "slaves.json"), slavesDir) + m := Machine{MachineID: "machine-1", ComputerName: "61414-PC"} + + if _, err := reg.Create(m, CreateInput{Folder: folder}); err != nil { + t.Fatalf("Create: %v", err) + } + for _, path := range []string{stateDir, slavesDir} { + info, err := os.Stat(path) + if err != nil { + t.Fatal(err) + } + if got := info.Mode().Perm(); got != 0o700 { + t.Fatalf("%s mode=%#o, want 0700", path, got) + } + } +} + func TestRegistryRejectsInvalidPathCharactersInName(t *testing.T) { dir := t.TempDir() folder := filepath.Join(dir, "repo") @@ -551,13 +730,6 @@ func TestRegistryRejectsInvalidPathCharactersInName(t *testing.T) { for _, name := range []string{ `bad\name`, "bad/name", - "bad:name", - "bad*name", - "bad?name", - `bad"name`, - "badname", - "bad|name", } { t.Run(name, func(t *testing.T) { reg := NewRegistry(filepath.Join(dir, name+".json"), filepath.Join(dir, "slaves", name)) diff --git a/internal/ui/orchestrator_real.go b/internal/ui/orchestrator_real.go index 50cc526..970d893 100644 --- a/internal/ui/orchestrator_real.go +++ b/internal/ui/orchestrator_real.go @@ -402,11 +402,11 @@ func (r *realOrchestrator) EnsureVSCode(ctx context.Context, ch chan<- ProgressE func (r *realOrchestrator) configureSharedCodex(ctx context.Context) error { _ = ctx - if err := codex.UpdateConfig(r.d.CodexConfigPath, codex.ModelserverProxySettings(modelproxy.DefaultBaseURL)); err != nil { + if err := codex.UpdateConfig(r.d.CodexConfigPath, codex.ModelserverProxySettings(modelproxy.DefaultBaseURL, codex.LegacyLocalProxyAPIKeyValue)); err != nil { return err } - _ = env.PersistUserEnv(codex.LocalProxyAPIKeyEnv, codex.LocalProxyAPIKeyValue) - _ = os.Setenv(codex.LocalProxyAPIKeyEnv, codex.LocalProxyAPIKeyValue) + _ = env.PersistUserEnv(codex.LocalProxyAPIKeyEnv, codex.LegacyLocalProxyAPIKeyValue) + _ = os.Setenv(codex.LocalProxyAPIKeyEnv, codex.LegacyLocalProxyAPIKeyValue) if r.d.TokenRefresherExePath != "" { _ = tokenrefresh.StartDaemon(r.d.TokenRefresherExePath) } @@ -728,7 +728,7 @@ func (r *realOrchestrator) LaunchAndShutdown(ctx context.Context) error { return fmt.Errorf("VS Code path unknown; was vscode_install completed?") } cmd := exec.Command(s.VSCode.Path, vscode.LaunchArgs(r.d.VSCodeUserDataDir, r.d.VSCodeExtDir)...) - cmd.Env = vscode.UpsertEnv(os.Environ(), codex.LocalProxyAPIKeyEnv, codex.LocalProxyAPIKeyValue) + cmd.Env = vscode.UpsertEnv(os.Environ(), codex.LocalProxyAPIKeyEnv, codex.LegacyLocalProxyAPIKeyValue) if r.d.TokenRefresherExePath != "" { _ = tokenrefresh.StartDaemon(r.d.TokenRefresherExePath) } diff --git a/internal/ui/orchestrator_real_test.go b/internal/ui/orchestrator_real_test.go index c41958e..786a9bf 100644 --- a/internal/ui/orchestrator_real_test.go +++ b/internal/ui/orchestrator_real_test.go @@ -224,7 +224,7 @@ func TestConfigureVSCodeSetsStableLocalCodexKey(t *testing.T) { if err := r.ConfigureVSCode(context.Background()); err != nil { t.Fatalf("configure: %v", err) } - if got := os.Getenv(codex.LocalProxyAPIKeyEnv); got != codex.LocalProxyAPIKeyValue { + if got := os.Getenv(codex.LocalProxyAPIKeyEnv); got != codex.LegacyLocalProxyAPIKeyValue { t.Fatalf("%s=%q, want stable local key", codex.LocalProxyAPIKeyEnv, got) } if got := os.Getenv("OPENAI_API_KEY"); got != "old-openai-token" { @@ -322,7 +322,7 @@ exit 0 if err != nil { t.Fatal(err) } - if strings.TrimSpace(string(envBody)) != codex.LocalProxyAPIKeyValue { + if strings.TrimSpace(string(envBody)) != codex.LegacyLocalProxyAPIKeyValue { t.Fatalf("%s child env = %q", codex.LocalProxyAPIKeyEnv, envBody) } argsBody, err := os.ReadFile(argsFile) @@ -1278,10 +1278,10 @@ func TestConfigureCodexDesktopWritesSharedConfigOnly(t *testing.T) { if !strings.Contains(string(b), `base_url = "`+modelproxy.DefaultBaseURL+`"`) { t.Fatalf("config missing local proxy base_url:\n%s", b) } - if !strings.Contains(string(b), `experimental_bearer_token = "`+codex.LocalProxyAPIKeyValue+`"`) { + if !strings.Contains(string(b), `experimental_bearer_token = "`+codex.LegacyLocalProxyAPIKeyValue+`"`) { t.Fatalf("config missing local proxy bearer token:\n%s", b) } - if got := os.Getenv(codex.LocalProxyAPIKeyEnv); got != codex.LocalProxyAPIKeyValue { + if got := os.Getenv(codex.LocalProxyAPIKeyEnv); got != codex.LegacyLocalProxyAPIKeyValue { t.Fatalf("%s=%q, want stable local key", codex.LocalProxyAPIKeyEnv, got) } if got := os.Getenv("OPENAI_API_KEY"); got != "old-openai-token" { diff --git a/internal/vscode/linux_package_test.go b/internal/vscode/linux_package_test.go index 9c53733..c85b94c 100644 --- a/internal/vscode/linux_package_test.go +++ b/internal/vscode/linux_package_test.go @@ -69,6 +69,42 @@ func TestPackageLinuxConsumesCrossLinuxBuildOutputs(t *testing.T) { } } +func TestPackageLinuxCreatesReproducibleArchives(t *testing.T) { + body, err := os.ReadFile("../../scripts/package-linux.sh") + if err != nil { + t.Fatal(err) + } + text := string(body) + for _, want := range []string{ + `--sort=name`, + `--owner=0`, + `--group=0`, + `--numeric-owner`, + `--mtime=@0`, + } { + if !strings.Contains(text, want) { + t.Fatalf("package-linux.sh missing reproducible tar option %q", want) + } + } +} + +func TestPackageLinuxSupportsDryRun(t *testing.T) { + body, err := os.ReadFile("../../scripts/package-linux.sh") + if err != nil { + t.Fatal(err) + } + text := string(body) + for _, want := range []string{ + `DRY_RUN="${DRY_RUN:-0}"`, + `if [[ "$DRY_RUN" == "1" ]]`, + `dry-run`, + } { + if !strings.Contains(text, want) { + t.Fatalf("package-linux.sh missing dry-run support %q", want) + } + } +} + func TestPackageLinuxPreflightsBinariesBeforeDownloads(t *testing.T) { body, err := os.ReadFile("../../scripts/package-linux.sh") if err != nil { @@ -91,6 +127,17 @@ func TestPackageLinuxPreflightsBinariesBeforeDownloads(t *testing.T) { } } +func TestLinuxPackageCommonRetriesAllCurlErrors(t *testing.T) { + body, err := os.ReadFile("../../scripts/linux-package-common.sh") + if err != nil { + t.Fatal(err) + } + text := string(body) + if !strings.Contains(text, `--retry-all-errors`) { + t.Fatalf("linux-package-common.sh should retry transient curl failures:\n%s", text) + } +} + func TestMakefileExposesLinuxTargets(t *testing.T) { body, err := os.ReadFile("../../Makefile") if err != nil { @@ -109,3 +156,14 @@ func TestMakefileExposesLinuxTargets(t *testing.T) { } } } + +func TestMakefileDisablesCGOForCrossLinux(t *testing.T) { + body, err := os.ReadFile("../../Makefile") + if err != nil { + t.Fatal(err) + } + text := string(body) + if !strings.Contains(text, `CGO_ENABLED=0 GOOS=linux GOARCH=$$arch`) { + t.Fatalf("Makefile cross-linux should disable CGO:\n%s", text) + } +} diff --git a/packaging/linux/codex-manifest-linux-amd64.json b/packaging/linux/codex-manifest-linux-amd64.json index 7d4784e..98b4f01 100644 --- a/packaging/linux/codex-manifest-linux-amd64.json +++ b/packaging/linux/codex-manifest-linux-amd64.json @@ -15,7 +15,8 @@ "shasum": "d272f0fa7bd1daf091c419fa22411362cc6b418e", "urls": [ "https://registry.npmmirror.com/@openai/codex/-/codex-0.139.0-linux-x64.tgz", - "https://npmreg.proxy.ustclug.org/@openai/codex/-/codex-0.139.0-linux-x64.tgz" + "https://npmreg.proxy.ustclug.org/@openai/codex/-/codex-0.139.0-linux-x64.tgz", + "https://registry.npmjs.org/@openai/codex/-/codex-0.139.0-linux-x64.tgz" ] } } diff --git a/packaging/linux/codex-manifest-linux-arm64.json b/packaging/linux/codex-manifest-linux-arm64.json index 08575c7..b1845cf 100644 --- a/packaging/linux/codex-manifest-linux-arm64.json +++ b/packaging/linux/codex-manifest-linux-arm64.json @@ -15,7 +15,8 @@ "shasum": "b94e071a7148e21e12d5deea1c6476255f96384b", "urls": [ "https://registry.npmmirror.com/@openai/codex/-/codex-0.139.0-linux-arm64.tgz", - "https://npmreg.proxy.ustclug.org/@openai/codex/-/codex-0.139.0-linux-arm64.tgz" + "https://npmreg.proxy.ustclug.org/@openai/codex/-/codex-0.139.0-linux-arm64.tgz", + "https://registry.npmjs.org/@openai/codex/-/codex-0.139.0-linux-arm64.tgz" ] } } diff --git a/scripts/linux-package-common.sh b/scripts/linux-package-common.sh index 5912287..93ad2fd 100644 --- a/scripts/linux-package-common.sh +++ b/scripts/linux-package-common.sh @@ -39,7 +39,7 @@ download_loom_asset() { rm -f "$cache" "$cache.part" echo "Fetching loom $asset ..." echo " URL: $url" - if ! curl --fail --location --retry 2 --retry-delay 2 --output "$cache.part" "$url"; then + if ! curl --fail --location --retry 2 --retry-all-errors --retry-delay 2 --output "$cache.part" "$url"; then rm -f "$cache.part" echo "ERROR: failed to download loom $asset" >&2 exit 2 diff --git a/scripts/package-linux.sh b/scripts/package-linux.sh index 36a2e64..d7e8940 100755 --- a/scripts/package-linux.sh +++ b/scripts/package-linux.sh @@ -6,6 +6,7 @@ cd "$ROOT" VERSION="${VERSION:-0.1.2}" OUT="${OUT:-dist}" +DRY_RUN="${DRY_RUN:-0}" source scripts/linux-package-common.sh mkdir -p "$OUT/cache/loom/$LOOM_RELEASE" "$OUT/linux" @@ -45,13 +46,23 @@ download_arch_assets() { } stage_arch() { - local arch stage tarball driver_asset driver_sha slave_asset slave_sha + local arch stage tarball driver_asset driver_sha slave_asset slave_sha arch="$1" stage="$OUT/linux/agentserver-linux-$arch" tarball="$OUT/linux/agentserver-linux-$arch.tar.gz" read -r driver_asset driver_sha < <(loom_asset_for_arch driver "$arch") read -r slave_asset slave_sha < <(loom_asset_for_arch slave "$arch") + if [[ "$DRY_RUN" == "1" ]]; then + echo "dry-run: would stage $stage" + echo "dry-run: agentserver <- $OUT/linux/$arch/agentserver" + echo "dry-run: driver-agent <- $(loom_cache_path "$driver_asset")" + echo "dry-run: slave-agent <- $(loom_cache_path "$slave_asset")" + echo "dry-run: manifest <- packaging/linux/codex-manifest-linux-$arch.json" + echo "dry-run: would write $tarball" + return 0 + fi + rm -rf "$stage" mkdir -p "$stage" cp "$OUT/linux/$arch/agentserver" "$stage/agentserver" @@ -64,7 +75,7 @@ stage_arch() { chmod 0755 "$stage/agentserver" "$stage/driver-agent" "$stage/slave-agent" rm -f "$tarball" "$tarball.sha256" - (cd "$OUT/linux" && tar -czf "agentserver-linux-$arch.tar.gz" "agentserver-linux-$arch") + (cd "$OUT/linux" && tar --sort=name --owner=0 --group=0 --numeric-owner --mtime=@0 -czf "agentserver-linux-$arch.tar.gz" "agentserver-linux-$arch") (cd "$OUT/linux" && sha256sum "agentserver-linux-$arch.tar.gz" > "agentserver-linux-$arch.tar.gz.sha256") echo "wrote $tarball" }