v0.2.1 (#55)

Co-authored-by: Tasnim Kabir Sadik <tksadik@omukk.dev>

Reviewed-on: wrenn/wrenn#55
Co-authored-by: pptx704 <rafeed@omukk.dev>
Co-committed-by: pptx704 <rafeed@omukk.dev>

internal/api/handlers_exec_stream.go

@@ -130,22 +130,8 @@ func (h *execStreamHandler) runExecStream(ctx context.Context, conn *websocket.C
 
 	// Forward stream events to WebSocket.
 	for stream.Receive() {
-		resp := stream.Msg()
-		switch ev := resp.Event.(type) {
-		case *pb.ExecStreamResponse_Start:
-			writeWSJSON(conn, wsOutMsg{Type: "start", PID: ev.Start.Pid})
-
-		case *pb.ExecStreamResponse_Data:
-			switch o := ev.Data.Output.(type) {
-			case *pb.ExecStreamData_Stdout:
-				writeWSJSON(conn, wsOutMsg{Type: "stdout", Data: string(o.Stdout)})
-			case *pb.ExecStreamData_Stderr:
-				writeWSJSON(conn, wsOutMsg{Type: "stderr", Data: string(o.Stderr)})
-			}
-
-		case *pb.ExecStreamResponse_End:
-			exitCode := ev.End.ExitCode
-			writeWSJSON(conn, wsOutMsg{Type: "exit", ExitCode: &exitCode})
+		if m, ok := procRespToWSMsg(stream.Msg()); ok {
+			writeWSJSON(conn, m)
 		}
 	}
 
@@ -159,6 +145,38 @@ func (h *execStreamHandler) runExecStream(ctx context.Context, conn *websocket.C
 	updateLastActive(h.db, sandboxID, sandboxIDStr)
 }
 
+// procStreamResp is satisfied by both *pb.ExecStreamResponse and
+// *pb.ConnectProcessResponse: their oneof events carry the same inner messages,
+// so the wire-to-WS mapping below is shared between the exec-stream and
+// connect-process handlers.
+type procStreamResp interface {
+	GetStart() *pb.ExecStreamStart
+	GetData() *pb.ExecStreamData
+	GetEnd() *pb.ExecStreamEnd
+}
+
+// procRespToWSMsg maps one process stream response to the WS message to send.
+// The bool is false when the response carries nothing to forward.
+func procRespToWSMsg(resp procStreamResp) (wsOutMsg, bool) {
+	if s := resp.GetStart(); s != nil {
+		return wsOutMsg{Type: "start", PID: s.Pid}, true
+	}
+	if d := resp.GetData(); d != nil {
+		switch o := d.Output.(type) {
+		case *pb.ExecStreamData_Stdout:
+			return wsOutMsg{Type: "stdout", Data: string(o.Stdout)}, true
+		case *pb.ExecStreamData_Stderr:
+			return wsOutMsg{Type: "stderr", Data: string(o.Stderr)}, true
+		}
+		return wsOutMsg{}, false
+	}
+	if e := resp.GetEnd(); e != nil {
+		exitCode := e.ExitCode
+		return wsOutMsg{Type: "exit", ExitCode: &exitCode}, true
+	}
+	return wsOutMsg{}, false
+}
+
 func sendWSError(conn *websocket.Conn, msg string) {
 	writeWSJSON(conn, wsOutMsg{Type: "error", Data: msg})
 }

internal/api/handlers_process.go

@@ -192,22 +192,8 @@ func (h *processHandler) runConnectProcess(ctx context.Context, conn *websocket.
 
 	// Forward stream events to WebSocket.
 	for stream.Receive() {
-		resp := stream.Msg()
-		switch ev := resp.Event.(type) {
-		case *pb.ConnectProcessResponse_Start:
-			writeWSJSON(conn, wsOutMsg{Type: "start", PID: ev.Start.Pid})
-
-		case *pb.ConnectProcessResponse_Data:
-			switch o := ev.Data.Output.(type) {
-			case *pb.ExecStreamData_Stdout:
-				writeWSJSON(conn, wsOutMsg{Type: "stdout", Data: string(o.Stdout)})
-			case *pb.ExecStreamData_Stderr:
-				writeWSJSON(conn, wsOutMsg{Type: "stderr", Data: string(o.Stderr)})
-			}
-
-		case *pb.ConnectProcessResponse_End:
-			exitCode := ev.End.ExitCode
-			writeWSJSON(conn, wsOutMsg{Type: "exit", ExitCode: &exitCode})
+		if m, ok := procRespToWSMsg(stream.Msg()); ok {
+			writeWSJSON(conn, m)
 		}
 	}
 

internal/api/middleware.go

@@ -60,6 +60,10 @@ func agentErrToHTTP(err error) (int, string, string) {
 		return http.StatusServiceUnavailable, "no_hosts_available", "no servers available — try again later"
 	case connect.CodeUnimplemented:
 		return http.StatusNotImplemented, "agent_error", err.Error()
+	case connect.CodeDeadlineExceeded:
+		return http.StatusGatewayTimeout, "timeout", "command timed out"
+	case connect.CodeInternal:
+		return http.StatusInternalServerError, "agent_error", err.Error()
 	default:
 		return http.StatusBadGateway, "agent_error", err.Error()
 	}

internal/api/sandbox_event_consumer.go

@@ -144,7 +144,7 @@ func (c *SandboxEventConsumer) handleMessage(ctx context.Context, msg redis.XMes
 		}
 	case events.CapsulePause:
 		if event.Outcome == events.OutcomeSuccess {
-			c.handleAutoPaused(ctx, sandboxID)
+			c.handleAutoPaused(ctx, sandboxID, event)
 		}
 	case events.CapsuleDestroy:
 		if event.Outcome == events.OutcomeSuccess {
@@ -226,12 +226,35 @@ func (c *SandboxEventConsumer) handleStarted(ctx context.Context, sandboxID pgty
 	}
 }
 
-func (c *SandboxEventConsumer) handleAutoPaused(ctx context.Context, sandboxID pgtype.UUID) {
+// handleAutoPaused reflects an autonomous (TTL reaper / shutdown) pause in the
+// DB and writes the audit row for it. The audit write happens only when the
+// status flip actually applied, so a stream redelivery does not double-count,
+// and so the HostMonitor host_state_sync fallback (which audits the
+// callback-lost case) stays mutually exclusive with this path.
+//
+// Uses audit.Log (row only) — NOT LogSandboxAutoPause, which republishes a
+// CapsulePause/system event that would loop straight back into this consumer.
+func (c *SandboxEventConsumer) handleAutoPaused(ctx context.Context, sandboxID pgtype.UUID, event events.Event) {
 	for _, fromStatus := range []string{"running", "pausing"} {
 		if _, err := c.db.UpdateSandboxStatusIf(ctx, db.UpdateSandboxStatusIfParams{
 			ID: sandboxID, Status: fromStatus, Status_2: "paused",
 		}); err == nil {
-			slog.Debug("sandbox event consumer: auto-paused fallback applied", "sandbox_id", id.FormatSandboxID(sandboxID), "from", fromStatus)
+			slog.Debug("sandbox event consumer: auto-paused applied", "sandbox_id", id.FormatSandboxID(sandboxID), "from", fromStatus)
+			reason := event.Metadata["reason"]
+			if reason == "" {
+				reason = "ttl_expired"
+			}
+			teamID, _ := id.ParseTeamID(event.TeamID)
+			c.audit.Log(ctx, audit.Entry{
+				TeamID:       teamID,
+				ActorType:    "system",
+				ResourceType: "sandbox",
+				ResourceID:   id.FormatSandboxID(sandboxID),
+				Action:       "pause",
+				Scope:        "team",
+				Status:       "info",
+				Metadata:     map[string]any{"reason": reason},
+			})
 			return
 		}
 	}

internal/api/sse_relay.go

@@ -104,6 +104,14 @@ func (r *SSERelay) handleMessage(ctx context.Context, msg *redis.Message) {
 		if err != nil {
 			slog.Debug("sse relay: sandbox hydration failed (may be deleted)", "sandbox_id", event.Resource.ID, "error", err)
 		} else {
+			// Override the hydrated status with the status implied by the event
+			// verb. Autonomous transitions (e.g. TTL auto-pause) flip the DB row
+			// in a separate stream consumer that races this Pub/Sub read, so the
+			// hydrated row may still carry the pre-transition status. The event
+			// itself is authoritative for the resulting state.
+			if status, ok := impliedSandboxStatus(event); ok {
+				sb.Status = status
+			}
 			payload.Sandbox = sb
 		}
 	}
@@ -138,6 +146,25 @@ func (r *SSERelay) hydrateSandbox(ctx context.Context, sandboxIDStr string) (*sa
 	return &resp, nil
 }
 
+// impliedSandboxStatus maps a successful capsule lifecycle event to the
+// sandbox status it results in. Used to override a hydrated DB row that may
+// still carry the pre-transition status because the reconciliation consumer
+// that flips it races this Pub/Sub read. Returns false for events with no
+// single deterministic resulting status (failures, destroy, state_changed).
+func impliedSandboxStatus(event events.Event) (string, bool) {
+	if event.Outcome != events.OutcomeSuccess {
+		return "", false
+	}
+	switch event.Event {
+	case events.CapsulePause:
+		return "paused", true
+	case events.CapsuleResume, events.CapsuleCreate:
+		return "running", true
+	default:
+		return "", false
+	}
+}
+
 func isCapsuleEvent(eventType string) bool {
 	switch eventType {
 	case events.CapsuleCreate, events.CapsulePause, events.CapsuleResume, events.CapsuleDestroy, events.CapsuleStateChanged: