fix(office): stop agent completion timestamps from staying blank after most runs (#2902)

Co-authored-by: nova28 <17953305+nova28@users.noreply.github.com>
Co-authored-by: Carlos Florencio <carlosmflorencio@gmail.com>

apps/backend/internal/office/repository/sqlite/runtime.go

@@ -67,13 +67,21 @@ func (r *Repository) ListAgentRuntimes(ctx context.Context) (map[string]*Runtime
 }
 
 // UpdateRuntimeLastRunFinished records the time an agent's run finished.
+// Upserts because UpsertAgentRuntime's only production caller is
+// onboarding, so an agent that never onboarded through that path has no
+// office_agent_runtime row yet — a bare UPDATE would silently affect 0
+// rows and the stamp would be lost. A newly-created row defaults to
+// status "idle" / empty pause_reason and must not clobber either on
+// conflict, matching UpsertAgentRuntime's own defaults.
 func (r *Repository) UpdateRuntimeLastRunFinished(ctx context.Context, agentID string, finishedAt time.Time) error {
 	now := time.Now().UTC()
 	_, err := r.db.ExecContext(ctx, r.db.Rebind(`
-		UPDATE office_agent_runtime
-		SET last_run_finished_at = ?, updated_at = ?
-		WHERE agent_id = ?
-	`), finishedAt, now, agentID)
+		INSERT INTO office_agent_runtime (agent_id, status, pause_reason, last_run_finished_at, updated_at)
+		VALUES (?, 'idle', '', ?, ?)
+		ON CONFLICT(agent_id) DO UPDATE SET
+			last_run_finished_at = excluded.last_run_finished_at,
+			updated_at = excluded.updated_at
+	`), agentID, finishedAt, now)
 	return err
 }
 

apps/backend/internal/office/repository/sqlite/runtime_test.go

@@ -115,7 +115,7 @@ func TestUpdateRuntimeLastRunFinished(t *testing.T) {
 	repo := newTestRepo(t)
 	ctx := context.Background()
 
-	if err := repo.UpsertAgentRuntime(ctx, "agent-1", "idle", ""); err != nil {
+	if err := repo.UpsertAgentRuntime(ctx, "agent-1", "paused", "budget"); err != nil {
 		t.Fatalf("upsert: %v", err)
 	}
 
@@ -135,4 +135,49 @@ func TestUpdateRuntimeLastRunFinished(t *testing.T) {
 	if diff < -time.Second || diff > time.Second {
 		t.Errorf("last_run_finished_at off by %v", diff)
 	}
+	if state.Status != "paused" {
+		t.Errorf("status = %q after UpdateRuntimeLastRunFinished, want paused", state.Status)
+	}
+	if state.PauseReason != "budget" {
+		t.Errorf("pause_reason = %q after UpdateRuntimeLastRunFinished, want budget", state.PauseReason)
+	}
+}
+
+// TestUpdateRuntimeLastRunFinished_CreatesRowWhenMissing is the regression
+// test for the "43 of 44 agents" defect: agents that never went through
+// onboarding's UpsertAgentRuntime call have no office_agent_runtime row, so
+// a bare UPDATE affects 0 rows and the stamp is silently lost. The write
+// must create the row instead of assuming one already exists.
+func TestUpdateRuntimeLastRunFinished_CreatesRowWhenMissing(t *testing.T) {
+	repo := newTestRepo(t)
+	ctx := context.Background()
+
+	if state, err := repo.GetAgentRuntime(ctx, "agent-no-row"); err != nil {
+		t.Fatalf("get before update: %v", err)
+	} else if state != nil {
+		t.Fatalf("expected no runtime row before update, got %+v", state)
+	}
+
+	now := time.Now().UTC()
+	if err := repo.UpdateRuntimeLastRunFinished(ctx, "agent-no-row", now); err != nil {
+		t.Fatalf("update: %v", err)
+	}
+
+	state, err := repo.GetAgentRuntime(ctx, "agent-no-row")
+	if err != nil {
+		t.Fatalf("get after update: %v", err)
+	}
+	if state == nil {
+		t.Fatal("expected the runtime row to be created")
+	}
+	if state.LastRunFinishedAt == nil {
+		t.Fatal("expected non-nil last_run_finished_at")
+	}
+	diff := state.LastRunFinishedAt.Sub(now)
+	if diff < -time.Second || diff > time.Second {
+		t.Errorf("last_run_finished_at off by %v", diff)
+	}
+	if state.Status != "idle" {
+		t.Errorf("status = %q, want idle default", state.Status)
+	}
 }

apps/backend/internal/office/service/event_subscribers.go

@@ -363,6 +363,7 @@ func (s *Service) handleAgentCompleted(ctx context.Context, event *bus.Event) er
 		return err
 	}
 	s.releaseTaskCheckoutForRun(ctx, run)
+	s.stampRunFinished(ctx, run)
 	return nil
 }
 
@@ -423,6 +424,7 @@ func (s *Service) handleTasklessAgentCompleted(
 		return err
 	}
 	s.releaseTaskCheckoutForRun(ctx, run)
+	s.stampRunFinished(ctx, run)
 	return nil
 }
 

apps/backend/internal/office/service/scheduler_checkout_release_test.go

@@ -3,6 +3,7 @@ package service_test
 import (
 	"context"
 	"testing"
+	"time"
 
 	"github.com/kandev/kandev/internal/common/logger"
 	"github.com/kandev/kandev/internal/events"
@@ -20,6 +21,14 @@ import (
 // checkout) never reached. Before the fix, handleAgentCompleted called
 // Service.FinishRun directly, which never released the checkout, so
 // every other agent was permanently blocked from acquiring the task.
+//
+// It also covers the sibling divergence in the same async path: the
+// SAME early return skips SchedulerIntegration.finishRun's
+// UpdateRuntimeLastRunFinished stamp, so office_agent_runtime.
+// last_run_finished_at is never written for any run that actually
+// launches an agent. Asserting only that the run reached "finished"
+// (the pre-existing assertion below) passes today and proves nothing
+// about either leak.
 func TestSchedulerTick_AgentCompletedReleasesTaskCheckout(t *testing.T) {
 	mock := &mockTaskStarter{}
 	svc := newTestService(t, service.ServiceOptions{TaskStarter: mock})
@@ -29,6 +38,7 @@ func TestSchedulerTick_AgentCompletedReleasesTaskCheckout(t *testing.T) {
 	if err := svc.RegisterEventSubscribers(eb); err != nil {
 		t.Fatalf("register subscribers: %v", err)
 	}
+	beforePublish := time.Now().UTC()
 
 	agent := &models.AgentInstance{
 		ID:                 "profile-checkout-release",
@@ -87,6 +97,17 @@ func TestSchedulerTick_AgentCompletedReleasesTaskCheckout(t *testing.T) {
 	if !ok {
 		t.Fatal("expected task checkout to be released after AgentCompleted, but another agent still cannot check it out")
 	}
+
+	runtime, err := svc.GetAgentRuntimeForTest(ctx, agent.ID)
+	if err != nil {
+		t.Fatalf("get agent runtime: %v", err)
+	}
+	if runtime == nil || runtime.LastRunFinishedAt == nil {
+		t.Fatal("expected last_run_finished_at to be stamped after AgentCompleted, but it is still unset")
+	}
+	if runtime.LastRunFinishedAt.Before(beforePublish) {
+		t.Errorf("last_run_finished_at = %v, want at/after %v", runtime.LastRunFinishedAt, beforePublish)
+	}
 }
 
 // TestSchedulerTick_AgentFailedReleasesTaskCheckout is the failure-path
@@ -161,6 +182,68 @@ func TestSchedulerTick_AgentFailedReleasesTaskCheckout(t *testing.T) {
 	}
 }
 
+// TestSchedulerTick_TasklessAgentCompletedStampsRuntime is the taskless-path
+// parity check for the same async completion divergence: handleAgentCompleted
+// routes a task_id-less event to handleTasklessAgentCompleted (the
+// agent_heartbeat-cron completion path), which must also stamp
+// office_agent_runtime.last_run_finished_at. The run is claimed directly via
+// ClaimNextRun rather than a full scheduler tick, because a taskless run's
+// own processRun→prepareAndLaunch path finishes it synchronously before any
+// launch (taskID == "" bypasses the "leave it claimed for the async
+// subscriber" branch) — going through the tick would finish the run before
+// the event ever had anything claimed to resolve.
+func TestSchedulerTick_TasklessAgentCompletedStampsRuntime(t *testing.T) {
+	svc := newTestService(t)
+	svc.SetSyncHandlers(true)
+	ctx := context.Background()
+	eb := bus.NewMemoryEventBus(logger.Default())
+	if err := svc.RegisterEventSubscribers(eb); err != nil {
+		t.Fatalf("register subscribers: %v", err)
+	}
+	beforePublish := time.Now().UTC()
+
+	agent := &models.AgentInstance{
+		ID:          "profile-taskless-stamp",
+		WorkspaceID: "ws-1",
+		Name:        "taskless-stamp-worker",
+		Role:        models.AgentRoleWorker,
+		Status:      models.AgentStatusIdle,
+	}
+	if err := svc.CreateAgentInstance(ctx, agent); err != nil {
+		t.Fatalf("create agent: %v", err)
+	}
+
+	if err := svc.QueueRun(ctx, agent.ID, service.RunReasonHeartbeat, `{}`, ""); err != nil {
+		t.Fatalf("queue: %v", err)
+	}
+	claimed, err := svc.ClaimNextRun(ctx)
+	if err != nil {
+		t.Fatalf("claim: %v", err)
+	}
+	if claimed == nil {
+		t.Fatal("expected a run to be claimed")
+	}
+
+	event := bus.NewEvent(events.AgentCompleted, "test", map[string]string{
+		"session_id": "session-taskless-stamp",
+		"agent_id":   agent.ID,
+	})
+	if err := eb.Publish(ctx, events.AgentCompleted, event); err != nil {
+		t.Fatalf("publish agent completed: %v", err)
+	}
+
+	runtime, err := svc.GetAgentRuntimeForTest(ctx, agent.ID)
+	if err != nil {
+		t.Fatalf("get agent runtime: %v", err)
+	}
+	if runtime == nil || runtime.LastRunFinishedAt == nil {
+		t.Fatal("expected last_run_finished_at to be stamped after taskless AgentCompleted, but it is still unset")
+	}
+	if runtime.LastRunFinishedAt.Before(beforePublish) {
+		t.Errorf("last_run_finished_at = %v, want at/after %v", runtime.LastRunFinishedAt, beforePublish)
+	}
+}
+
 // TestSchedulerTick_ReapsStaleCheckoutWithNoInFlightRun pins the tick-loop
 // wiring for the stale-checkout reaper (backstop for
 // releaseTaskCheckoutForRun): a checkout with nothing behind it must be

apps/backend/internal/office/service/scheduler_integration.go

@@ -719,11 +719,7 @@ func (si *SchedulerIntegration) finishRun(
 	}
 
 	si.releaseCheckoutIfNeeded(ctx, run)
-
-	// Record cooldown timestamp in-memory and DB.
-	now := time.Now().UTC()
-	agent.LastRunFinishedAt = &now
-	_ = si.svc.repo.UpdateRuntimeLastRunFinished(ctx, agent.ID, now)
+	si.svc.stampRunFinished(ctx, run)
 
 	si.svc.LogActivityWithRun(ctx, agent.WorkspaceID,
 		"scheduler", "office-scheduler",

apps/backend/internal/office/service/scheduler_runs.go

@@ -4,6 +4,7 @@ import (
 	"context"
 	"database/sql"
 	"errors"
+	"time"
 
 	"go.uber.org/zap"
 
@@ -110,6 +111,25 @@ func (s *Service) releaseTaskCheckoutForRun(ctx context.Context, run *models.Run
 	}
 }
 
+// stampRunFinished records the cooldown timestamp for the agent that ran
+// this run. Shared by the synchronous SchedulerIntegration.finishRun path
+// and the async AgentCompleted/AgentStopped event subscribers so the two
+// completion paths cannot drift apart again — see releaseTaskCheckoutForRun
+// for the same split. run.AgentProfileID is the launching agent's own
+// identity (processRun resolves the agent from the run), so no extra
+// agent load is needed here.
+func (s *Service) stampRunFinished(ctx context.Context, run *models.Run) {
+	if run == nil || run.AgentProfileID == "" {
+		return
+	}
+	if err := s.repo.UpdateRuntimeLastRunFinished(ctx, run.AgentProfileID, time.Now().UTC()); err != nil {
+		s.logger.Error("failed to stamp agent runtime last_run_finished_at",
+			zap.String("run_id", run.ID),
+			zap.String("agent_id", run.AgentProfileID),
+			zap.Error(err))
+	}
+}
+
 // publishRunProcessed emits an OfficeRunProcessed bus event with the
 // per-run context the WS gateway needs to fan the update out.
 // Best-effort: skips silently when no bus is configured.

apps/backend/internal/office/service/test_helpers.go

@@ -87,6 +87,12 @@ func (s *Service) GetTaskAssigneeForTest(ctx context.Context, taskID string) (st
 	return s.repo.GetTaskAssignee(ctx, taskID)
 }
 
+// GetAgentRuntimeForTest exposes the repo's runtime lookup for service
+// package tests asserting the last_run_finished_at stamp.
+func (s *Service) GetAgentRuntimeForTest(ctx context.Context, agentID string) (*sqlite.RuntimeState, error) {
+	return s.repo.GetAgentRuntime(ctx, agentID)
+}
+
 // RunSchedulerTick runs a single scheduler tick for testing.
 // This exercises the full processRun pipeline including task launch.
 func RunSchedulerTick(svc *Service, ctx context.Context) {