From 276ab9dcf23089a12200aa2e731c1ab3d48483ca Mon Sep 17 00:00:00 2001 From: Yanpeng Wang Date: Mon, 27 Jul 2026 19:23:54 +0800 Subject: [PATCH] feat: scope sandboxes to sessions --- README.md | 12 +- docs/architecture.md | 7 +- docs/architecture/runtime-and-sandbox.md | 38 ++- docs/compatibility.md | 3 +- docs/roadmap.md | 7 +- internal/app/sandbox_test_helpers_test.go | 60 +++++ internal/app/session_sandbox_test.go | 301 ++++++++++++++++++++++ internal/app/session_service.go | 53 ++-- internal/app/session_service_test.go | 2 + internal/sandbox/session.go | 96 +++++++ internal/sandbox/session_test.go | 289 +++++++++++++++++++++ 11 files changed, 838 insertions(+), 30 deletions(-) create mode 100644 internal/app/sandbox_test_helpers_test.go create mode 100644 internal/app/session_sandbox_test.go create mode 100644 internal/sandbox/session.go create mode 100644 internal/sandbox/session_test.go diff --git a/README.md b/README.md index 64305b6..d240835 100644 --- a/README.md +++ b/README.md @@ -49,8 +49,9 @@ compatibility is not complete. | Storage | SQLite; one process; at-least-once restart recovery | Important gaps include interrupt propagation, durable per-step runtime output, -session-scoped workspaces, retries with side-effect idempotency, MCP execution, -files/skills/memory, resolved multiagent orchestration, and distributed workers. +durable sandbox checkpoint/restore across process restart, retries with +side-effect idempotency, MCP execution, files/skills/memory, resolved multiagent +orchestration, and distributed workers. See the [roadmap](docs/roadmap.md). ## Quick start @@ -119,8 +120,11 @@ go run ./cmd/managed-agent serve Docker sandboxes use `--network none` by default, but containers still share the host kernel and this path has not been audited for hostile multi-tenant -workloads. Sandboxes are currently provisioned per run, so filesystem state -does not persist across session turns. +workloads. Sandboxes are scoped to the session: the first run needing tools +provisions one and later runs in the same session reuse it, so filesystem state +persists across turns; the sandbox is released when the session is deleted. The +manager is in-memory, so a process restart does not restore an idle session's +sandbox. ## Documentation diff --git a/docs/architecture.md b/docs/architecture.md index 1c9bef8..4d7cdfb 100644 --- a/docs/architecture.md +++ b/docs/architecture.md @@ -116,8 +116,11 @@ The strongest current risks are semantic rather than structural: effect without a durable journal of the prior attempt. 3. Pending client actions are encoded in events and stop reasons rather than a first-class durable `pending_actions` model. -4. Sandboxes are per run rather than per session, so workspace continuity is - not yet part of the session contract. +4. Sandboxes are session-scoped: a session's logical sandbox is provisioned on + first tool use, reused across its runs, and released on session deletion. + The manager is in-memory, so a process restart does not restore an idle + session's workspace, and there is no durable checkpoint, quota, or eviction + policy yet. 5. `SessionService` currently combines session CRUD, admission, dispatch, and completion orchestration. These responsibilities should be separated before introducing multiple workers or richer retry behavior. diff --git a/docs/architecture/runtime-and-sandbox.md b/docs/architecture/runtime-and-sandbox.md index 6ecbcbd..efe2f32 100644 --- a/docs/architecture/runtime-and-sandbox.md +++ b/docs/architecture/runtime-and-sandbox.md @@ -91,12 +91,38 @@ Containers share the host kernel. This provider has not been audited for hostile multi-tenant use; stronger isolation such as gVisor or a remote sandbox can be added behind the same provider interface. -## Current lifecycle limitation - -Sandboxes are provisioned per run and destroyed when the run ends. Files -created by tools do not persist across session turns. Moving to a -session-scoped workspace requires an explicit ownership, checkpoint, quota, and -cleanup model rather than simply retaining temporary directories. +## Session-scoped ownership + +A sandbox is scoped to the session, not to a single run. The first run in a +session that needs tools provisions a logical sandbox; every later run in the +same session reuses that same instance, so filesystem state a tool creates in +one run is visible to the next. Different sessions acquire under different keys +and never share a sandbox, so they stay isolated even when they use the same +agent and environment. + +Ownership lives in a session-scoped manager that wraps the provider inside the +`internal/sandbox` package: acquisition provisions on first use and returns the +cached instance afterwards; release destroys it. The `AgentRuntime` is unaware +of this — the application resolves the sandbox and passes it in the run request. + +Entering idle does not tear the sandbox down; it stays live between turns. +Deleting the session releases it, running the provider teardown exactly once. A +provisioning failure is not cached, so a later run may retry. + +The manager holds sandboxes in memory. Restart does not restore an idle +session's sandbox: a process restart starts from an empty workspace, and the +first run after restart provisions a fresh one. Durable checkpoint/restore is +not implemented in this slice. Quotas and eviction are also out of scope here. + +This is a process-boundary limitation, not just a persistence gap. Because +ownership lives only in the in-memory manager, a new process cannot reattach to +sandboxes an earlier process provisioned. A crash or an ungraceful restart +therefore leaves those provider resources — Docker containers or local temp +directories — orphaned, since the only code that would tear them down (`Release` +on session deletion) died with the process. Nothing reclaims them until an +external cleanup step or a reaper exists, and neither is built yet. Reclaiming +in-flight sandboxes on shutdown (a shutdown manager or reaper) is out of scope +for this slice. ## Streaming previews diff --git a/docs/compatibility.md b/docs/compatibility.md index 0a34c50..683e934 100644 --- a/docs/compatibility.md +++ b/docs/compatibility.md @@ -74,7 +74,8 @@ tests in `internal/httpapi/sdk_golden_test.go`. | Custom-tool handoff (`agent.custom_tool_use` → park → `user.custom_tool_result` → resume) | partial | `TestWorkshop_CustomToolHandoff`, `TestSessionService_CustomToolParksAndResumes`, `TestAgentCore_CustomToolParksWithRequiresAction`, `TestFake_CustomToolResultEndsIdle`, `TestProjectMessages_CustomToolResultPairing` | A custom tool call parks the run at `session.status_idle` with `stop_reason.type == "requires_action"` and `event_ids` naming the committed `agent.custom_tool_use`; a `user.custom_tool_result` referencing that id resumes a fresh run to `end_turn`. End-to-end proven; the exact `requires_action` payload shape is unconfirmed against the official wire. | | Built-in tool run end-to-end via the app layer | partial | `TestSessionService_BuiltinToolRunEndToEnd` | A session-driven run provisions a sandbox, executes a built-in tool, and reaches `end_turn` through the durable run path. | | `always_ask` permission policy on built-in tools | partial | `TestAgentCore_AlwaysAskBuiltinParks` | An `always_ask` built-in call parks the run (`agent.tool_use{evaluated_permission:"ask"}` + `requires_action`); the **resume** path (`user.tool_confirmation` → projected `tool_result` + built-in execution) is not wired. `ProjectMessages` drops the dangling `tool_use` so the parked call never poisons a real request (`TestProjectMessages_DropsDanglingToolUse`, `TestProjectMessages_DropsDanglingCustomToolUse`). | -| Local sandbox (`Provider`/`Sandbox`, restricted local process) | partial | `TestLocal_ExecEcho`, `TestLocal_FileRoundTripAndConfinement`, `TestLocal_Timeout` | `internal/sandbox` provides a two-layer interface and a local-process default that confines paths to a work dir, clears the environment, applies a timeout, and caps output. **Dev-grade guardrail, not a security boundary — do not run untrusted code.** The sandbox is **per-run**: it is provisioned at the start of each run and destroyed when the run ends, so tool-produced file state does **not** persist across turns. Session-scoped persistence remains a later slice behind the same interface. | +| Session-scoped sandbox lifecycle (`internal/sandbox` `SessionManager`) | partial | `TestSessionManager_ReusesSandboxPerSession`, `TestSessionManager_IsolatesSessions`, `TestSessionManager_ReleaseDestroysExactlyOnce`, `TestSessionService_SandboxPersistsAcrossRuns`, `TestSessionService_SandboxIsolatedBetweenSessions`, `TestSessionService_SandboxProvisionedOncePerSession`, `TestSessionService_IdleDoesNotDestroySandbox`, `TestSessionService_DeleteReleasesSandboxExactlyOnce` | A sandbox is **scoped to the session**, not the run. The first run needing tools provisions one logical sandbox; later runs in the same session reuse it, so tool-produced file state **persists across turns**. Different sessions get distinct sandboxes and stay isolated. Entering idle does not destroy the sandbox; deleting the session releases it exactly once. Ownership lives in a session-scoped manager that wraps the provider inside the `sandbox` package; `AgentRuntime` receives a resolved sandbox and is unaware of the lifecycle. The manager is in-memory: a process restart does not restore an idle session's sandbox, and there is no durable checkpoint, quota, or eviction yet. | +| Local sandbox (`Provider`/`Sandbox`, restricted local process) | partial | `TestLocal_ExecEcho`, `TestLocal_FileRoundTripAndConfinement`, `TestLocal_Timeout` | `internal/sandbox` provides a two-layer interface and a local-process default that confines paths to a work dir, clears the environment, applies a timeout, and caps output. **Dev-grade guardrail, not a security boundary — do not run untrusted code.** Sandboxes are session-scoped (see the row above): provisioned on first tool use and reused across the session's runs. | | Docker sandbox (real-isolation `Provider`, opt-in) | partial | `TestDocker_*` (skipped without a daemon), `TestResolveSandboxProvider_DefaultsToLocal` | The same `Provider`/`Sandbox` interface has a Docker-backed implementation (shells out to the `docker` CLI, no extra module dependency). It gives **real isolation**: each sandbox is a container with its own Linux namespaces/cgroups, a separate filesystem, and `--network none` by default. Selected at startup via `MANAGED_AGENT_SANDBOX=docker` (default is local); image via `MANAGED_AGENT_SANDBOX_IMAGE` (defaults `alpine:latest`). gVisor (`--runtime=runsc`) can layer under the same interface later with no interface change. Not audited for hostile multi-tenant use (shared host kernel). Docker tests are gated on a running daemon and skip in default offline CI. | | MCP toolsets | unsupported | `TestParseTools_BuiltinCustomMCP` | Parsed into `domain.ToolSet` but never resolved or executed. | diff --git a/docs/roadmap.md b/docs/roadmap.md index f85a546..fdb387c 100644 --- a/docs/roadmap.md +++ b/docs/roadmap.md @@ -34,8 +34,8 @@ surface-area compatibility. ## 3. Session continuity -- Move from per-run to session-scoped workspaces with explicit lifecycle and - cleanup policies. +- Add durable checkpoint/restore so an idle session's sandbox survives a process + restart, plus quota and eviction policies for session-scoped workspaces. - Add context compaction, token usage accounting, and model-request spans. - Persist resumable runtime checkpoints where the public contract requires continuity. @@ -68,6 +68,9 @@ When real deployment requirements demand it: - Atomic input/run admission and single-node restart recovery. - Multi-step model/tool loop. - Local sandbox plus optional Docker provider. +- Session-scoped sandbox ownership: reused across a session's runs, isolated + between sessions, released on session deletion (in-memory manager; no durable + restore yet). - `bash`, `read`, `write`, `edit`, `glob`, and `grep` execution. - Custom-tool handoff with `requires_action`. - Opt-in streaming preview of `agent.message`. diff --git a/internal/app/sandbox_test_helpers_test.go b/internal/app/sandbox_test_helpers_test.go new file mode 100644 index 0000000..88a4ed4 --- /dev/null +++ b/internal/app/sandbox_test_helpers_test.go @@ -0,0 +1,60 @@ +package app + +import ( + "context" + "sync/atomic" + + "github.com/yanpgwang/managed-agent-go/internal/sandbox" +) + +// provisionCountingProvider wraps a Provider and counts Provision calls so a +// test can assert a session provisions its logical sandbox exactly once across +// repeated runs. +type provisionCountingProvider struct { + inner sandbox.Provider + provisions atomic.Int64 +} + +func (p *provisionCountingProvider) Provision(ctx context.Context, spec sandbox.Spec) (sandbox.Sandbox, error) { + p.provisions.Add(1) + return p.inner.Provision(ctx, spec) +} + +func (p *provisionCountingProvider) count() int64 { return p.provisions.Load() } + +// destroyCountingProvider hands out sandboxes that count their own Destroy calls +// so a test can assert session deletion tears the sandbox down exactly once. +type destroyCountingProvider struct { + inner sandbox.Provider + destroys atomic.Int64 +} + +func (p *destroyCountingProvider) Provision(ctx context.Context, spec sandbox.Spec) (sandbox.Sandbox, error) { + box, err := p.inner.Provision(ctx, spec) + if err != nil { + return nil, err + } + return &destroyCountingSandbox{inner: box, provider: p}, nil +} + +func (p *destroyCountingProvider) destroyCount() int64 { return p.destroys.Load() } + +type destroyCountingSandbox struct { + inner sandbox.Sandbox + provider *destroyCountingProvider +} + +func (s *destroyCountingSandbox) Exec(ctx context.Context, cmd sandbox.Command) (*sandbox.Result, error) { + return s.inner.Exec(ctx, cmd) +} +func (s *destroyCountingSandbox) ReadFile(ctx context.Context, path string) ([]byte, error) { + return s.inner.ReadFile(ctx, path) +} +func (s *destroyCountingSandbox) WriteFile(ctx context.Context, path string, data []byte) error { + return s.inner.WriteFile(ctx, path, data) +} +func (s *destroyCountingSandbox) Root() string { return s.inner.Root() } +func (s *destroyCountingSandbox) Destroy(ctx context.Context) error { + s.provider.destroys.Add(1) + return s.inner.Destroy(ctx) +} diff --git a/internal/app/session_sandbox_test.go b/internal/app/session_sandbox_test.go new file mode 100644 index 0000000..403c1f9 --- /dev/null +++ b/internal/app/session_sandbox_test.go @@ -0,0 +1,301 @@ +package app + +import ( + "context" + "strings" + "testing" + "time" + + "github.com/yanpgwang/managed-agent-go/internal/agentruntime" + "github.com/yanpgwang/managed-agent-go/internal/domain" + "github.com/yanpgwang/managed-agent-go/internal/sandbox" + "github.com/yanpgwang/managed-agent-go/internal/store" +) + +// fileToolRuntime is a test AgentRuntime that exercises the request sandbox +// directly. It interprets the trigger's user text as a tiny script: +// +// "write " writes data to path in the sandbox, then +// "read " reads path and emits its contents as an agent.message. +// +// It lets a test prove that filesystem state a tool creates in one run is (or is +// not) visible to a later run through the sandbox the app hands the runtime. +type fileToolRuntime struct{} + +func (fileToolRuntime) Run( + ctx context.Context, + req agentruntime.RunRequest, + sink agentruntime.EventSink, +) (agentruntime.RunOutcome, error) { + text := triggerText(req.Trigger) + fields := strings.Fields(text) + reply := "" + if len(fields) >= 2 && req.Sandbox != nil { + switch fields[0] { + case "write": + data := "" + if len(fields) >= 3 { + data = strings.Join(fields[2:], " ") + } + if err := req.Sandbox.WriteFile(ctx, fields[1], []byte(data)); err != nil { + reply = "write-error: " + err.Error() + } else { + reply = "wrote " + fields[1] + } + case "read": + data, err := req.Sandbox.ReadFile(ctx, fields[1]) + if err != nil { + reply = "read-error: " + err.Error() + } else { + reply = "read: " + string(data) + } + } + } + _, err := sink.Emit(ctx, []domain.EventDraft{{ + Type: domain.EvAgentMessage, + Payload: map[string]any{"content": []any{map[string]any{"type": "text", "text": reply}}}, + }}) + return agentruntime.RunOutcome{}, err +} + +func triggerText(ev domain.Event) string { + if text, ok := ev.Payload["text"].(string); ok { + return text + } + blocks, _ := ev.Payload["content"].([]any) + for _, b := range blocks { + if m, ok := b.(map[string]any); ok { + if t, _ := m["text"].(string); t != "" { + return t + } + } + } + return "" +} + +// toolAgent creates an agent whose toolset is non-empty so the session provisions +// a sandbox, plus a cloud environment. Returns their ids. +func toolAgentAndEnv(t *testing.T, as *AgentService, envs *EnvironmentService) (string, string) { + t.Helper() + ctx := context.Background() + ag, err := as.Create(ctx, domain.Agent{ + Name: "tool-agent", + Model: domain.Model{ID: "claude-opus-4-8"}, + Tools: []any{map[string]any{"type": domain.BuiltinToolsetType}}, + }) + if err != nil { + t.Fatal(err) + } + env, err := envs.Create(ctx, domain.Environment{Name: "e", ConfigType: "cloud"}) + if err != nil { + t.Fatal(err) + } + return ag.ID, env.ID +} + +// lastAgentText returns the text of the most recent agent.message in history. +func lastAgentText(t *testing.T, ss *SessionService, sessionID string) string { + t.Helper() + hist, err := ss.events.History(context.Background(), sessionID, 0, 100000) + if err != nil { + t.Fatal(err) + } + var text string + for _, e := range hist { + if e.Type == domain.EvAgentMessage { + text = contentBlockText(e.Payload["content"]) + } + } + return text +} + +func newFileToolService(t *testing.T) (*SessionService, *AgentService, *EnvironmentService, *sandbox.SessionManager) { + t.Helper() + db, _ := store.OpenMemory() + t.Cleanup(func() { db.Close() }) + ids := domain.NewSeqIDGen() + clk := domain.FixedClock{T: time.Unix(1, 0).UTC()} + es := NewEventService(store.NewEventStore(db, ids, clk), NewHub(64)) + as := NewAgentService(store.NewAgentRepo(db), ids, clk) + envs := NewEnvironmentService(store.NewEnvironmentRepo(db), ids, clk) + ss := NewSessionService(store.NewSessionRepo(db), store.NewAgentRepo(db), store.NewEnvironmentRepo(db), + es, store.NewRunStore(db, ids, clk), fileToolRuntime{}, sandbox.NewLocalProvider(), ids, clk) + return ss, as, envs, ss.sandbox +} + +// TestSessionService_SandboxPersistsAcrossRuns proves the core session-scoped +// property: a file written by the first run is readable by a later run in the +// same session. +func TestSessionService_SandboxPersistsAcrossRuns(t *testing.T) { + ss, as, envs, _ := newFileToolService(t) + ctx := context.Background() + agID, envID := toolAgentAndEnv(t, as, envs) + + sess, err := ss.Create(ctx, CreateSessionInput{AgentID: agID, EnvironmentID: envID}) + if err != nil { + t.Fatal(err) + } + // Release the session's sandbox so its temp dir does not leak; Release is + // idempotent, so a later Delete in the test is still fine. + t.Cleanup(func() { _ = ss.sandbox.Release(context.Background(), sess.ID) }) + + if _, err := ss.SendEvent(ctx, sess.ID, []domain.EventDraft{{ + Type: domain.EvUserMessage, Payload: map[string]any{"text": "write state.txt hello-across-runs"}, + }}); err != nil { + t.Fatal(err) + } + pollUntilStatus(t, ss, sess.ID, domain.StatusIdle) + + if _, err := ss.SendEvent(ctx, sess.ID, []domain.EventDraft{{ + Type: domain.EvUserMessage, Payload: map[string]any{"text": "read state.txt"}, + }}); err != nil { + t.Fatal(err) + } + pollUntilStatus(t, ss, sess.ID, domain.StatusIdle) + + if got := lastAgentText(t, ss, sess.ID); got != "read: hello-across-runs" { + t.Fatalf("second run read %q, want the file written by the first run", got) + } +} + +// TestSessionService_SandboxIsolatedBetweenSessions proves a file written in one +// session's sandbox is not visible from a different session, even when both use +// the same agent and environment. +func TestSessionService_SandboxIsolatedBetweenSessions(t *testing.T) { + ss, as, envs, _ := newFileToolService(t) + ctx := context.Background() + agID, envID := toolAgentAndEnv(t, as, envs) + + writer, err := ss.Create(ctx, CreateSessionInput{AgentID: agID, EnvironmentID: envID}) + if err != nil { + t.Fatal(err) + } + t.Cleanup(func() { _ = ss.sandbox.Release(context.Background(), writer.ID) }) + if _, err := ss.SendEvent(ctx, writer.ID, []domain.EventDraft{{ + Type: domain.EvUserMessage, Payload: map[string]any{"text": "write secret.txt only-in-writer"}, + }}); err != nil { + t.Fatal(err) + } + pollUntilStatus(t, ss, writer.ID, domain.StatusIdle) + + other, err := ss.Create(ctx, CreateSessionInput{AgentID: agID, EnvironmentID: envID}) + if err != nil { + t.Fatal(err) + } + t.Cleanup(func() { _ = ss.sandbox.Release(context.Background(), other.ID) }) + if _, err := ss.SendEvent(ctx, other.ID, []domain.EventDraft{{ + Type: domain.EvUserMessage, Payload: map[string]any{"text": "read secret.txt"}, + }}); err != nil { + t.Fatal(err) + } + pollUntilStatus(t, ss, other.ID, domain.StatusIdle) + + got := lastAgentText(t, ss, other.ID) + if !strings.HasPrefix(got, "read-error:") { + t.Fatalf("different session saw %q, want a read error (isolated sandbox)", got) + } +} + +// TestSessionService_SandboxProvisionedOncePerSession proves repeated runs in one +// session reuse a single logical sandbox instead of provisioning a new one each +// run. +func TestSessionService_SandboxProvisionedOncePerSession(t *testing.T) { + db, _ := store.OpenMemory() + t.Cleanup(func() { db.Close() }) + ids := domain.NewSeqIDGen() + clk := domain.FixedClock{T: time.Unix(1, 0).UTC()} + es := NewEventService(store.NewEventStore(db, ids, clk), NewHub(64)) + as := NewAgentService(store.NewAgentRepo(db), ids, clk) + envs := NewEnvironmentService(store.NewEnvironmentRepo(db), ids, clk) + counting := &provisionCountingProvider{inner: sandbox.NewLocalProvider()} + ss := NewSessionService(store.NewSessionRepo(db), store.NewAgentRepo(db), store.NewEnvironmentRepo(db), + es, store.NewRunStore(db, ids, clk), fileToolRuntime{}, counting, ids, clk) + + ctx := context.Background() + agID, envID := toolAgentAndEnv(t, as, envs) + sess, err := ss.Create(ctx, CreateSessionInput{AgentID: agID, EnvironmentID: envID}) + if err != nil { + t.Fatal(err) + } + t.Cleanup(func() { _ = ss.sandbox.Release(context.Background(), sess.ID) }) + + for i := 0; i < 3; i++ { + if _, err := ss.SendEvent(ctx, sess.ID, []domain.EventDraft{{ + Type: domain.EvUserMessage, Payload: map[string]any{"text": "write f.txt data"}, + }}); err != nil { + t.Fatal(err) + } + pollUntilStatus(t, ss, sess.ID, domain.StatusIdle) + } + + if got := counting.count(); got != 1 { + t.Fatalf("provisioned %d sandboxes across 3 runs, want 1", got) + } +} + +// TestSessionService_IdleDoesNotDestroySandbox proves that reaching idle after a +// run leaves the session's sandbox intact (still holding its files), rather than +// destroying it. +func TestSessionService_IdleDoesNotDestroySandbox(t *testing.T) { + ss, as, envs, mgr := newFileToolService(t) + ctx := context.Background() + agID, envID := toolAgentAndEnv(t, as, envs) + sess, err := ss.Create(ctx, CreateSessionInput{AgentID: agID, EnvironmentID: envID}) + if err != nil { + t.Fatal(err) + } + t.Cleanup(func() { _ = mgr.Release(context.Background(), sess.ID) }) + + if _, err := ss.SendEvent(ctx, sess.ID, []domain.EventDraft{{ + Type: domain.EvUserMessage, Payload: map[string]any{"text": "write keep.txt still-here"}, + }}); err != nil { + t.Fatal(err) + } + pollUntilStatus(t, ss, sess.ID, domain.StatusIdle) + + // After idle the sandbox must still exist and hold the file: acquiring it + // again returns the same live instance rather than a fresh empty one. + box, err := mgr.Acquire(ctx, sess.ID, sandbox.Spec{}) + if err != nil { + t.Fatal(err) + } + data, err := box.ReadFile(ctx, "keep.txt") + if err != nil || string(data) != "still-here" { + t.Fatalf("sandbox lost its file after idle: data=%q err=%v", data, err) + } +} + +// TestSessionService_DeleteReleasesSandboxExactlyOnce proves deleting a session +// tears its sandbox down exactly once. +func TestSessionService_DeleteReleasesSandboxExactlyOnce(t *testing.T) { + db, _ := store.OpenMemory() + t.Cleanup(func() { db.Close() }) + ids := domain.NewSeqIDGen() + clk := domain.FixedClock{T: time.Unix(1, 0).UTC()} + es := NewEventService(store.NewEventStore(db, ids, clk), NewHub(64)) + as := NewAgentService(store.NewAgentRepo(db), ids, clk) + envs := NewEnvironmentService(store.NewEnvironmentRepo(db), ids, clk) + counting := &destroyCountingProvider{inner: sandbox.NewLocalProvider()} + ss := NewSessionService(store.NewSessionRepo(db), store.NewAgentRepo(db), store.NewEnvironmentRepo(db), + es, store.NewRunStore(db, ids, clk), fileToolRuntime{}, counting, ids, clk) + + ctx := context.Background() + agID, envID := toolAgentAndEnv(t, as, envs) + sess, err := ss.Create(ctx, CreateSessionInput{AgentID: agID, EnvironmentID: envID}) + if err != nil { + t.Fatal(err) + } + if _, err := ss.SendEvent(ctx, sess.ID, []domain.EventDraft{{ + Type: domain.EvUserMessage, Payload: map[string]any{"text": "write f.txt data"}, + }}); err != nil { + t.Fatal(err) + } + pollUntilStatus(t, ss, sess.ID, domain.StatusIdle) + + if err := ss.Delete(ctx, sess.ID); err != nil { + t.Fatal(err) + } + if got := counting.destroyCount(); got != 1 { + t.Fatalf("session delete destroyed the sandbox %d times, want 1", got) + } +} diff --git a/internal/app/session_service.go b/internal/app/session_service.go index 1e86cbb..ff2c400 100644 --- a/internal/app/session_service.go +++ b/internal/app/session_service.go @@ -68,7 +68,7 @@ type SessionService struct { events *EventService runs *store.RunStore rt agentruntime.AgentRuntime - sandbox sandbox.Provider + sandbox *sandbox.SessionManager ids domain.IDGenerator clock domain.Clock @@ -81,7 +81,7 @@ func NewSessionService(sess *store.SessionRepo, agents *store.AgentRepo, envs *s sandboxProvider sandbox.Provider, ids domain.IDGenerator, clock domain.Clock, ) *SessionService { return &SessionService{sess: sess, agents: agents, envs: envs, events: events, runs: runs, rt: rt, - sandbox: sandboxProvider, ids: ids, clock: clock, lockSeed: maphash.MakeSeed()} + sandbox: sandbox.NewSessionManager(sandboxProvider), ids: ids, clock: clock, lockSeed: maphash.MakeSeed()} } func (s *SessionService) lockFor(id string) *sync.Mutex { @@ -252,18 +252,37 @@ func (s *SessionService) Archive(ctx context.Context, id string) (domain.Session func (s *SessionService) Delete(ctx context.Context, id string) error { lock := s.lockFor(id) lock.Lock() - defer lock.Unlock() sess, err := s.sess.Get(ctx, id) if err != nil { + lock.Unlock() return err } if sess.Status == domain.StatusRunning { + lock.Unlock() return domain.Conflict("cannot delete a running session; interrupt first") } if err := s.sess.Delete(ctx, id); err != nil { + lock.Unlock() return err } + // The durable delete is committed under the shard lock. Drop the lock before + // the external provider teardown and stream close: sandbox Destroy is an + // out-of-process call (a Docker container or host temp dir) that must not + // hold the shard — and thus stall every other session hashing to it — for + // its whole duration. Each path above unlocks exactly once, so there is no + // double unlock. + lock.Unlock() + + // Permanently clean up the session's logical sandbox. Release is idempotent + // and a no-op when the session never provisioned one; when it did, it runs + // the provider teardown exactly once. This is an external call made after the + // durable delete, outside the store transaction and outside the shard lock. + if err := s.sandbox.Release(ctx, id); err != nil { + // A teardown failure can leak sandbox resources (a host temp dir or a + // container). Log it; the session is already durably deleted. + log.Printf("delete: sandbox release failed session_id=%s: %v", id, err) + } now := s.clock.Now().UTC() s.events.CloseSession(id, domain.Event{ @@ -309,18 +328,22 @@ func (s *SessionService) drainRuns(sessionID string) { runErr = toolErr } - // Provision one sandbox per run when the session has tools, and destroy - // it when the run ends. Provisioning and destruction are external calls - // made outside any transaction or process lock. A provisioning failure - // terminates the session like any other unrecoverable runtime error. + // Resolve the session's logical sandbox when the session has tools. The + // sandbox is scoped to the session, not this run: the first run that needs + // tools provisions it and later runs in the same session reuse the same + // instance, so filesystem state a tool created in an earlier run is visible + // now. Acquisition is an external call made outside any transaction or + // process lock. A provisioning failure terminates the session like any + // other unrecoverable runtime error. The sandbox is NOT destroyed when the + // run ends; it is released only when the session is deleted (see Delete). var box sandbox.Sandbox if runErr == nil && toolSetHasTools(toolSet) { - box, err = s.sandbox.Provision(ctx, sandbox.Spec{Timeout: 120 * time.Second}) + box, err = s.sandbox.Acquire(ctx, sessionID, sandbox.Spec{Timeout: 120 * time.Second}) if err != nil { - log.Printf("drain: sandbox provision failed session_id=%s run_id=%s: %v", sessionID, runID, err) + log.Printf("drain: sandbox acquire failed session_id=%s run_id=%s: %v", sessionID, runID, err) runErr = err } else { - log.Printf("drain: sandbox provisioned session_id=%s run_id=%s", sessionID, runID) + log.Printf("drain: sandbox acquired session_id=%s run_id=%s", sessionID, runID) } } @@ -367,11 +390,11 @@ func (s *SessionService) drainRuns(sessionID string) { } if box != nil { - if destroyErr := box.Destroy(ctx); destroyErr != nil { - // Destroy failure can leak sandbox resources (host temp dir or a - // container). Log it; the run outcome itself is unaffected. - log.Printf("drain: sandbox destroy failed session_id=%s run_id=%s: %v", sessionID, runID, destroyErr) - } + // The sandbox is session-scoped: it deliberately outlives this run so + // the next run in the session sees the filesystem state this run left + // behind. Teardown happens on session deletion (Delete releases it), + // not here. + log.Printf("drain: sandbox retained for session session_id=%s run_id=%s", sessionID, runID) } drafts := sink.Drafts() diff --git a/internal/app/session_service_test.go b/internal/app/session_service_test.go index 85ba22e..b44ad98 100644 --- a/internal/app/session_service_test.go +++ b/internal/app/session_service_test.go @@ -1038,6 +1038,7 @@ func TestSessionService_BuiltinToolRunEndToEnd(t *testing.T) { if err != nil { t.Fatal(err) } + t.Cleanup(func() { _ = sessions.sandbox.Release(context.Background(), session.ID) }) pollUntilStatus(t, sessions, session.ID, domain.StatusIdle) @@ -1104,6 +1105,7 @@ func TestSessionService_CustomToolParksAndResumes(t *testing.T) { if err != nil { t.Fatal(err) } + t.Cleanup(func() { _ = sessions.sandbox.Release(context.Background(), session.ID) }) // The run parks: session goes idle awaiting the custom tool result. pollUntilStatus(t, sessions, session.ID, domain.StatusIdle) diff --git a/internal/sandbox/session.go b/internal/sandbox/session.go new file mode 100644 index 0000000..c070082 --- /dev/null +++ b/internal/sandbox/session.go @@ -0,0 +1,96 @@ +package sandbox + +import ( + "context" + "sync" +) + +// SessionManager gives each session a single logical sandbox that lives across +// its runs. The first run that needs tools acquires the sandbox; later runs in +// the same session get the same instance back, so filesystem state a tool +// created in one run is visible to the next. Different sessions acquire under +// different keys and never share a sandbox. +// +// Provider-specific provisioning and teardown stay behind this type: it wraps a +// Provider and owns the acquire/reuse/release lifecycle the application needs, +// without the application (or AgentRuntime) knowing how a sandbox is created or +// destroyed. Entering idle does nothing here; only Release tears a sandbox down, +// and it does so exactly once. +type SessionManager struct { + provider Provider + + mu sync.Mutex + boxes map[string]*managedSandbox +} + +// managedSandbox is one session's cached sandbox. ready is closed once +// provisioning finishes, so concurrent Acquire callers for the same session +// block until the single Provision call completes rather than racing to +// provision twice. +type managedSandbox struct { + ready chan struct{} + box Sandbox + err error +} + +// NewSessionManager wraps a Provider so sandboxes can be scoped to a session +// instead of a single run. +func NewSessionManager(provider Provider) *SessionManager { + return &SessionManager{provider: provider, boxes: make(map[string]*managedSandbox)} +} + +// Acquire returns the session's sandbox, provisioning it via the wrapped +// provider on first use and returning the same instance on every later call for +// the same sessionID. Concurrent first-use callers provision exactly once and +// all receive the same sandbox. A provisioning failure is not cached: the entry +// is dropped so a later Acquire can try again. +// +// spec is used only for the initial provision; once a session has a sandbox its +// spec is fixed and later specs are ignored (the existing instance is reused). +func (m *SessionManager) Acquire(ctx context.Context, sessionID string, spec Spec) (Sandbox, error) { + m.mu.Lock() + if entry, ok := m.boxes[sessionID]; ok { + m.mu.Unlock() + <-entry.ready + return entry.box, entry.err + } + entry := &managedSandbox{ready: make(chan struct{})} + m.boxes[sessionID] = entry + m.mu.Unlock() + + entry.box, entry.err = m.provider.Provision(ctx, spec) + if entry.err != nil { + // Do not cache a failed provision: drop the entry so a subsequent run's + // Acquire re-attempts rather than reusing the failure forever. + m.mu.Lock() + if m.boxes[sessionID] == entry { + delete(m.boxes, sessionID) + } + m.mu.Unlock() + } + close(entry.ready) + return entry.box, entry.err +} + +// Release permanently destroys the session's sandbox and forgets it. Deleting +// the map entry under the lock guarantees the underlying Destroy runs at most +// once even under concurrent Release calls, so a session's sandbox is cleaned +// up exactly once. Releasing an unknown or already-released session is a no-op. +func (m *SessionManager) Release(ctx context.Context, sessionID string) error { + m.mu.Lock() + entry, ok := m.boxes[sessionID] + if ok { + delete(m.boxes, sessionID) + } + m.mu.Unlock() + if !ok { + return nil + } + // Wait for any in-flight provision so we destroy the real instance rather + // than racing a half-provisioned entry. + <-entry.ready + if entry.err != nil || entry.box == nil { + return nil + } + return entry.box.Destroy(ctx) +} diff --git a/internal/sandbox/session_test.go b/internal/sandbox/session_test.go new file mode 100644 index 0000000..f48cecf --- /dev/null +++ b/internal/sandbox/session_test.go @@ -0,0 +1,289 @@ +package sandbox + +import ( + "context" + "errors" + "sync" + "sync/atomic" + "testing" + "time" +) + +// countingProvider wraps a Provider and counts Provision calls so tests can +// assert a session provisions its logical sandbox exactly once. +type countingProvider struct { + inner Provider + provisions atomic.Int64 + provisionErr error +} + +func (p *countingProvider) Provision(ctx context.Context, spec Spec) (Sandbox, error) { + p.provisions.Add(1) + if p.provisionErr != nil { + return nil, p.provisionErr + } + return p.inner.Provision(ctx, spec) +} + +// countingSandbox wraps a Sandbox and counts Destroy calls so tests can assert +// a session's sandbox is torn down exactly once. +type countingSandbox struct { + inner Sandbox + destroys atomic.Int64 +} + +func (s *countingSandbox) Exec(ctx context.Context, cmd Command) (*Result, error) { + return s.inner.Exec(ctx, cmd) +} +func (s *countingSandbox) ReadFile(ctx context.Context, path string) ([]byte, error) { + return s.inner.ReadFile(ctx, path) +} +func (s *countingSandbox) WriteFile(ctx context.Context, path string, data []byte) error { + return s.inner.WriteFile(ctx, path, data) +} +func (s *countingSandbox) Root() string { return s.inner.Root() } +func (s *countingSandbox) Destroy(ctx context.Context) error { + s.destroys.Add(1) + return s.inner.Destroy(ctx) +} + +// destroyCountingProvider hands out countingSandbox instances and remembers the +// last one it created so a test can inspect its destroy count. +type destroyCountingProvider struct { + inner Provider + mu sync.Mutex + last *countingSandbox +} + +func (p *destroyCountingProvider) Provision(ctx context.Context, spec Spec) (Sandbox, error) { + box, err := p.inner.Provision(ctx, spec) + if err != nil { + return nil, err + } + cs := &countingSandbox{inner: box} + p.mu.Lock() + p.last = cs + p.mu.Unlock() + return cs, nil +} + +func TestSessionManager_ReusesSandboxPerSession(t *testing.T) { + cp := &countingProvider{inner: NewLocalProvider()} + m := NewSessionManager(cp) + ctx := context.Background() + + first, err := m.Acquire(ctx, "sesn_a", Spec{}) + if err != nil { + t.Fatal(err) + } + second, err := m.Acquire(ctx, "sesn_a", Spec{}) + if err != nil { + t.Fatal(err) + } + if first != second { + t.Fatal("Acquire returned a different sandbox for the same session") + } + if got := cp.provisions.Load(); got != 1 { + t.Fatalf("provisions = %d, want 1 (session reuses one logical sandbox)", got) + } + if err := m.Release(ctx, "sesn_a"); err != nil { + t.Fatal(err) + } +} + +func TestSessionManager_IsolatesSessions(t *testing.T) { + m := NewSessionManager(NewLocalProvider()) + ctx := context.Background() + + a, err := m.Acquire(ctx, "sesn_a", Spec{}) + if err != nil { + t.Fatal(err) + } + b, err := m.Acquire(ctx, "sesn_b", Spec{}) + if err != nil { + t.Fatal(err) + } + if a == b || a.Root() == b.Root() { + t.Fatal("different sessions must get distinct sandboxes") + } + + if err := a.WriteFile(ctx, "shared.txt", []byte("A")); err != nil { + t.Fatal(err) + } + if _, err := b.ReadFile(ctx, "shared.txt"); err == nil { + t.Fatal("session B must not see a file written in session A") + } + _ = m.Release(ctx, "sesn_a") + _ = m.Release(ctx, "sesn_b") +} + +func TestSessionManager_ReleaseDestroysExactlyOnce(t *testing.T) { + dp := &destroyCountingProvider{inner: NewLocalProvider()} + m := NewSessionManager(dp) + ctx := context.Background() + + if _, err := m.Acquire(ctx, "sesn_a", Spec{}); err != nil { + t.Fatal(err) + } + box := dp.last + + if err := m.Release(ctx, "sesn_a"); err != nil { + t.Fatal(err) + } + // A second Release for the same (now forgotten) session must be a no-op. + if err := m.Release(ctx, "sesn_a"); err != nil { + t.Fatal(err) + } + // Releasing a session that never provisioned is also a no-op. + if err := m.Release(ctx, "sesn_never"); err != nil { + t.Fatal(err) + } + if got := box.destroys.Load(); got != 1 { + t.Fatalf("destroys = %d, want exactly 1", got) + } +} + +func TestSessionManager_ConcurrentAcquireProvisionsOnce(t *testing.T) { + cp := &countingProvider{inner: NewLocalProvider()} + m := NewSessionManager(cp) + ctx := context.Background() + + const goroutines = 32 + var wg sync.WaitGroup + boxes := make([]Sandbox, goroutines) + for i := 0; i < goroutines; i++ { + wg.Add(1) + go func(i int) { + defer wg.Done() + box, err := m.Acquire(ctx, "sesn_race", Spec{}) + if err != nil { + t.Errorf("acquire: %v", err) + return + } + boxes[i] = box + }(i) + } + wg.Wait() + + if got := cp.provisions.Load(); got != 1 { + t.Fatalf("concurrent Acquire provisioned %d times, want 1", got) + } + for i := 1; i < goroutines; i++ { + if boxes[i] != boxes[0] { + t.Fatal("concurrent Acquire returned different sandbox instances") + } + } + _ = m.Release(ctx, "sesn_race") +} + +func TestSessionManager_ProvisionFailureIsNotCached(t *testing.T) { + cp := &countingProvider{inner: NewLocalProvider(), provisionErr: errors.New("boom")} + m := NewSessionManager(cp) + ctx := context.Background() + + if _, err := m.Acquire(ctx, "sesn_a", Spec{}); err == nil { + t.Fatal("expected provision error") + } + // The failure must not be cached: a later Acquire should retry provisioning. + cp.provisionErr = nil + if _, err := m.Acquire(ctx, "sesn_a", Spec{}); err != nil { + t.Fatalf("retry after failed provision: %v", err) + } + if got := cp.provisions.Load(); got != 2 { + t.Fatalf("provisions = %d, want 2 (failure retried, not cached)", got) + } + _ = m.Release(ctx, "sesn_a") +} + +// blockingProvider makes Provision block until proceed is closed, so a test can +// drive Release into the exact window where a provision is still in flight. It +// signals entry on entered and records the countingSandbox it eventually hands +// out so the test can assert the real instance is destroyed. +type blockingProvider struct { + inner Provider + entered chan struct{} // closed once Provision has started + proceed chan struct{} // Provision blocks until this is closed + mu sync.Mutex + last *countingSandbox +} + +func (p *blockingProvider) Provision(ctx context.Context, spec Spec) (Sandbox, error) { + close(p.entered) + <-p.proceed + box, err := p.inner.Provision(ctx, spec) + if err != nil { + return nil, err + } + cs := &countingSandbox{inner: box} + p.mu.Lock() + p.last = cs + p.mu.Unlock() + return cs, nil +} + +// TestSessionManager_ReleaseWaitsForInflightProvision guards the narrow window +// where Release races a still-provisioning Acquire. With a deliberately blocking +// provider, Release begins while Provision is blocked: it must not return early, +// and once the provision finishes it must destroy the real instance exactly +// once rather than a half-provisioned entry. +func TestSessionManager_ReleaseWaitsForInflightProvision(t *testing.T) { + p := &blockingProvider{ + inner: NewLocalProvider(), + entered: make(chan struct{}), + proceed: make(chan struct{}), + } + m := NewSessionManager(p) + ctx := context.Background() + + acquired := make(chan struct{}) + go func() { + if _, err := m.Acquire(ctx, "sesn_a", Spec{}); err != nil { + t.Errorf("acquire: %v", err) + } + close(acquired) + }() + + // Wait until Provision is actually in flight (Acquire is blocked inside the + // provider), so Release runs against a still-provisioning entry. + select { + case <-p.entered: + case <-time.After(2 * time.Second): + t.Fatal("Provision was never entered") + } + + // Release now races the in-flight provision. It must NOT return until the + // provision finishes, because it has to destroy the real instance. + released := make(chan error, 1) + go func() { released <- m.Release(ctx, "sesn_a") }() + + select { + case <-released: + t.Fatal("Release returned before the in-flight provision completed") + case <-time.After(100 * time.Millisecond): + // Expected: Release is blocked waiting on the provision to finish. + } + + // Unblock the provision. Acquire completes and Release can now proceed to + // destroy the freshly provisioned sandbox. + close(p.proceed) + + select { + case err := <-released: + if err != nil { + t.Fatalf("release: %v", err) + } + case <-time.After(2 * time.Second): + t.Fatal("Release did not complete after the provision finished") + } + <-acquired + + p.mu.Lock() + box := p.last + p.mu.Unlock() + if box == nil { + t.Fatal("provider never handed out a sandbox") + } + if got := box.destroys.Load(); got != 1 { + t.Fatalf("destroys = %d, want exactly 1", got) + } +}