90b1a2a · crush · git.iko.soy

git.iko.soy/crush

feat(server): make server prompts independent of client connections

Give each workspace its own run lifetime and shutdown gate so agent work is not tied to the HTTP request that submitted it. Final state can still be saved during shutdown, avoiding lost completion or error records when a workspace is canceled. Co-Authored-By: Charm Crush <crush@charm.land>

commit90b1a2a67bf70a2f5fc1f998ccd84317cab5d68b authorChristian Rocha <christian@rocha.is> dateSun 31 May 2026, 21:48 (5 months ago) parentcd592e8

2 files changed+78−8

compare with main

unifiedsplit

internal/agent/agent.go+22−7 internal/backend/backend.go+56−1

internal/agent/agent.go+22−7

@@ -305,7 +305,13 @@ func (a *sessionAgent) Run(ctx context.Context, call SessionAgentCall) (result *

305305 // (the pubsub broker fan-in does not serialize publishes from

306306 // different upstream brokers).

307307 defer func() {

308− if flushErr := a.messages.FlushAll(ctx); flushErr != nil {

308+ // Use a context detached from the run context: workspace

309+ // shutdown cancels ctx before this goroutine returns, but the

310+ // buffered streaming deltas must still land before the DB is

311+ // closed. A short timeout bounds the flush.

312+ flushCtx, flushCancel := context.WithTimeout(context.WithoutCancel(ctx), 5*time.Second)

313+ defer flushCancel()

314+ if flushErr := a.messages.FlushAll(flushCtx); flushErr != nil {

309315 slog.Error("Failed to flush pending message updates after run", "error", flushErr)

310316 }

311317 if skipRunComplete {

@@ -577,11 +583,20 @@ func (a *sessionAgent) Run(ctx context.Context, call SessionAgentCall) (result *

577583 if currentAssistant == nil {

578584 return result, err

579585 }

586+ // Persist final state with a context detached from the run

587+ // context. The run context (ctx) is derived from the

588+ // workspace context, which workspace shutdown cancels before

589+ // agent goroutines finish; using ctx here would drop the

590+ // final assistant state. WithoutCancel keeps the values

591+ // (e.g. session ID) while ignoring cancellation, and a short

592+ // timeout bounds the cleanup writes.

593+ cleanupCtx, cleanupCancel := context.WithTimeout(context.WithoutCancel(ctx), 5*time.Second)

594+ defer cleanupCancel()

580595 // Ensure we finish thinking on error to close the reasoning state.

581596 currentAssistant.FinishThinking()

582597 toolCalls := currentAssistant.ToolCalls()

583− // INFO: we use the parent context here because the genCtx has been cancelled.

584− msgs, createErr := a.messages.List(ctx, currentAssistant.SessionID)

598+ // INFO: we use the cleanup context here because the genCtx has been cancelled.

599+ msgs, createErr := a.messages.List(cleanupCtx, currentAssistant.SessionID)

585600 if createErr != nil {

586601 return nil, createErr

587602 }

@@ -590,7 +605,7 @@ func (a *sessionAgent) Run(ctx context.Context, call SessionAgentCall) (result *

590605 tc.Finished = true

591606 tc.Input = "{}"

592607 currentAssistant.AddToolCall(tc)

593− updateErr := a.messages.Update(ctx, *currentAssistant)

608+ updateErr := a.messages.Update(cleanupCtx, *currentAssistant)

594609 if updateErr != nil {

595610 return nil, updateErr

596611 }

@@ -623,7 +638,7 @@ func (a *sessionAgent) Run(ctx context.Context, call SessionAgentCall) (result *

623638 Content: content,

624639 IsError: true,

625640 }

626− _, createErr = a.messages.Create(ctx, currentAssistant.SessionID, message.CreateMessageParams{

641+ _, createErr = a.messages.Create(cleanupCtx, currentAssistant.SessionID, message.CreateMessageParams{

627642 Role: message.Tool,

628643 Parts: []message.ContentPart{

629644 toolResult,

@@ -670,9 +685,9 @@ func (a *sessionAgent) Run(ctx context.Context, call SessionAgentCall) (result *

670685 } else {

671686 currentAssistant.AddFinish(message.FinishReasonError, defaultTitle, err.Error())

672687 }

673− // Note: we use the parent context here because the genCtx has been

688+ // Note: we use the cleanup context here because the genCtx has been

674689 // cancelled.

675− updateErr := a.messages.Update(ctx, *currentAssistant)

690+ updateErr := a.messages.Update(cleanupCtx, *currentAssistant)

676691 if updateErr != nil {

677692 return nil, updateErr

678693 }

internal/backend/backend.go+56−1

@@ -34,6 +34,7 @@ var (

3434 ErrUnknownCommand = errors.New("unknown command")

3535 ErrInvalidClientID = errors.New("invalid client_id")

3636 ErrClientNotAttached = errors.New("client not attached")

37+ ErrWorkspaceClosing = errors.New("workspace closing")

3738)

3839

3940// DefaultCreateGrace is the window in which a client must open an SSE

@@ -108,6 +109,23 @@ type Workspace struct {

108109 // with fallback to the cleaned absolute path.

109110 resolvedPath string

110111

112+ // ctx is the workspace-scoped run context. It is derived from

113+ // the backend context in CreateWorkspace and lives for the

114+ // lifetime of the workspace; cancel tears it down. Agent runs

115+ // dispatched on behalf of this workspace are bound to ctx so

116+ // their lifetime is owned by the workspace, not by any single

117+ // client's HTTP request.

118+ ctx context.Context

119+ cancel context.CancelFunc

120+

121+ // runMu guards closing and gates dispatch of new agent runs.

122+ // closing is set by Shutdown so no new runs are accepted once

123+ // teardown has begun. runWG tracks dispatched agent goroutines

124+ // so Shutdown can wait for them to return before app cleanup.

125+ runMu sync.Mutex

126+ closing bool

127+ runWG sync.WaitGroup

128+

111129 // clientsMu guards clients. It is held only briefly (no IO).

112130 clientsMu sync.Mutex

113131 // clients tracks each client's claim on this workspace. Refcount

@@ -122,7 +140,7 @@ type Workspace struct {

122140}

123141

124142// invokeShutdown calls the workspace shutdown hook if set, falling

125−// back to the embedded [app.App.Shutdown] when not.

143+// back to the workspace [Workspace.Shutdown] wrapper when not.

126144func (w *Workspace) invokeShutdown() {

127145 if w.shutdownFn != nil {

128146 w.shutdownFn()

@@ -133,6 +151,40 @@ func (w *Workspace) invokeShutdown() {

133151 }

134152}

135153

154+// Shutdown tears the workspace down in an order that is safe for

155+// agent runs whose lifetime is bound to the workspace context. It

156+// shadows the promoted [app.App.Shutdown] so callers reaching

157+// ws.Shutdown() always observe this ordering:

158+//

159+// 1. Mark the workspace closing so no new agent runs are accepted.

160+// 2. Cancel the workspace run context so any dispatched goroutine

161+// that has not yet registered its per-session cancel still

162+// observes cancellation.

163+// 3. Cancel active coordinator work for runs that already

164+// registered their per-session cancel function.

165+// 4. Wait for dispatched agent goroutines to return.

166+// 5. Run the embedded [app.App.Shutdown] cleanup (DB, LSP, etc).

167+//

168+// CancelAll is idempotent, so the second call inside app.App.Shutdown

169+// is harmless; the important guarantee is that cancel -> CancelAll ->

170+// runWG.Wait completes before the embedded cleanup touches the DB.

171+func (w *Workspace) Shutdown() {

172+ w.runMu.Lock()

173+ w.closing = true

174+ w.runMu.Unlock()

175+

176+ if w.cancel != nil {

177+ w.cancel()

178+ }

179+ if w.App != nil && w.AgentCoordinator != nil {

180+ w.AgentCoordinator.CancelAll()

181+ }

182+ w.runWG.Wait()

183+ if w.App != nil {

184+ w.App.Shutdown()

185+ }

186+}

187+

136188// New creates a new [Backend].

137189func New(ctx context.Context, cfg *config.ConfigStore, shutdownFn ShutdownFunc) *Backend {

138190 return &Backend{

@@ -247,6 +299,7 @@ func (b *Backend) CreateWorkspace(args proto.Workspace) (*Workspace, proto.Works

247299 return nil, proto.Workspace{}, fmt.Errorf("failed to create app workspace: %w", err)

248300 }

249301

302+ wsCtx, wsCancel := context.WithCancel(b.ctx)

250303 ws := &Workspace{

251304 App: appWorkspace,

252305 ID: id,

@@ -255,6 +308,8 @@ func (b *Backend) CreateWorkspace(args proto.Workspace) (*Workspace, proto.Works

255308 Env: args.Env,

256309 Skills: skillsMgr,

257310 resolvedPath: key,

311+ ctx: wsCtx,

312+ cancel: wsCancel,

258313 clients: make(map[string]*clientState),

259314 }

260315