refactor(aio): apply simplify-pass cleanups to telemetry and session state

Findings from a 4-angle simplify pass (reuse / simplification /
efficiency / altitude), no behavior change except where noted:

- Extract resolveEffectiveCloudStatus: the 9-line terminal-wins status
  precedence chain was duplicated byte-for-byte in deriveCloudRunState
  and deriveSessionViewState; three review angles flagged it
  independently. One exported helper in cloudRunState.ts now owns the
  policy.
- isTaskDetailNotFoundError uses requestErrorStatus(error) === 404
  instead of string-matching "[404]" in the message - the fetcher keeps
  that message format only as legacy, and the substring check could
  false-positive on a non-404 error quoting an upstream body. Behavior
  tightened accordingly (plain Errors mentioning [404] no longer read as
  404s).
- taskDetailQuery no longer retries 404s: the detail fetch is now
  always-on, and optimistic/cloud-pending tasks 404 by design, so
  react-query's default 3 retries turned every mount of a fresh task
  into 4 requests.
- pickFreshestTask is binary (all six call sites were), which also
  deletes the logically dead ?? fallbacks at the two hook call sites.
- sessionService: merge the duplicated optimistic-clear blocks and
  skip-apply early returns in performCloudTaskSessionHydration; shrink
  finalizeTerminalCloudTask to the (taskRunId, status) both call sites
  use; clearTerminalCloudPromptState derives taskId from the session
  (removing the dual-key lookup hazard) and now also serves the
  pre-existing inline terminal cleanup in handleCloudTaskUpdate, so the
  composer-settling invariant lives in one place; staleness is computed
  once per update in handleCloudTaskUpdate.
- OTel modules: normalizeMethod (the safety-relevant __posthog/
  double-prefix rule), asString, and the batch flush-interval constant
  move to otel-attributes.ts instead of being duplicated across
  otel-telemetry.ts and otel-trace-builder.ts.

Generated-By: PostHog Code
Task-Id: 0c511836-2180-455a-9b58-45df2a0661ec

Author: tatoalo

packages/agent/src/otel-attributes.ts

@@ -6,10 +6,30 @@ const MAX_ATTR_CHARS = 200;
 // session cleanup (the sandbox can be torn down right after), so keep it short.
 const EXPORT_TIMEOUT_MS = 5000;
 
+// Batch flush cadence shared by the log and span processors.
+const DEFAULT_FLUSH_INTERVAL_MS = 2000;
+
 export type AttributeValue = string | number | boolean;
 export type Attributes = Record<string, AttributeValue>;
 
-export { EXPORT_TIMEOUT_MS, MAX_ATTR_CHARS, MAX_BODY_CHARS };
+export {
+  DEFAULT_FLUSH_INTERVAL_MS,
+  EXPORT_TIMEOUT_MS,
+  MAX_ATTR_CHARS,
+  MAX_BODY_CHARS,
+};
+
+/**
+ * extNotification() can double-prefix custom methods (see matchesExt in
+ * acp-extensions.ts); normalize so both spellings map identically.
+ */
+export function normalizeMethod(method: string): string {
+  return method.startsWith("__posthog/") ? method.slice(1) : method;
+}
+
+export function asString(value: unknown): string | undefined {
+  return typeof value === "string" ? value : undefined;
+}
 
 export function truncate(value: string, max: number): string {
   return value.length <= max ? value : `${value.slice(0, max)}…`;

packages/agent/src/otel-telemetry.ts

@@ -13,9 +13,12 @@ import { POSTHOG_NOTIFICATIONS } from "./acp-extensions";
 import {
   type Attributes,
   asRecord,
+  asString,
+  DEFAULT_FLUSH_INTERVAL_MS,
   EXPORT_TIMEOUT_MS,
   entryTime,
   MAX_BODY_CHARS,
+  normalizeMethod,
   strAttr,
   truncate,
   usageAttributes,
@@ -26,7 +29,6 @@ import type { StoredNotification } from "./types";
 import type { Logger } from "./utils/logger";
 
 const SERVICE_NAME = "posthog-code-agent";
-const DEFAULT_FLUSH_INTERVAL_MS = 2000;
 
 export interface OtelTelemetryConfig {
   /** Full OTLP logs endpoint URL, e.g. "https://us.i.posthog.com/i/v1/logs" */
@@ -146,11 +148,7 @@ export function mapNotificationToLogRecord(
 ): MappedLogRecord | null {
   const rawMethod = entry.notification.method;
   if (typeof rawMethod !== "string") return null;
-  // extNotification() can double-prefix custom methods (see matchesExt in
-  // acp-extensions.ts); normalize so both spellings map identically.
-  const method = rawMethod.startsWith("__posthog/")
-    ? rawMethod.slice(1)
-    : rawMethod;
+  const method = normalizeMethod(rawMethod);
   const params = asRecord(entry.notification.params) ?? {};
 
   if (method === "session/update") {
@@ -212,7 +210,7 @@ export function mapNotificationToLogRecord(
       strAttr(attrs, "progress_group", params.group);
       const step = strAttr(attrs, "progress_step", params.step);
       const status = strAttr(attrs, "progress_status", params.status);
-      const label = typeof params.label === "string" ? params.label : undefined;
+      const label = asString(params.label);
       return record(
         INFO,
         `progress: ${step ?? "step"} ${status ?? ""}${label ? ` (${label})` : ""}`.trim(),

packages/agent/src/otel-trace-builder.ts

@@ -17,15 +17,16 @@ import { POSTHOG_NOTIFICATIONS } from "./acp-extensions";
 import {
   type Attributes,
   asRecord,
+  asString,
+  DEFAULT_FLUSH_INTERVAL_MS,
   EXPORT_TIMEOUT_MS,
   entryTime,
+  normalizeMethod,
   strAttr,
   usageAttributes,
 } from "./otel-attributes";
 import type { StoredNotification } from "./types";
 
-const DEFAULT_FLUSH_INTERVAL_MS = 2000;
-
 export interface RunTraceBuilderConfig {
   /** Full OTLP traces endpoint URL, e.g. "https://us.i.posthog.com/i/v1/traces" */
   url: string;
@@ -98,9 +99,7 @@ export class RunTraceBuilder {
     if (this.ended) return this.rootContext;
     const rawMethod = entry.notification.method;
     if (typeof rawMethod !== "string") return this.currentContext();
-    const method = rawMethod.startsWith("__posthog/")
-      ? rawMethod.slice(1)
-      : rawMethod;
+    const method = normalizeMethod(rawMethod);
     const params = asRecord(entry.notification.params) ?? {};
     const time = entryTime(entry.timestamp);
 
@@ -181,8 +180,7 @@ export class RunTraceBuilder {
 
   private endTurn(params: Record<string, unknown>, time: Date): Context {
     const context = this.turnContext ?? this.rootContext;
-    const stopReason =
-      typeof params.stopReason === "string" ? params.stopReason : undefined;
+    const stopReason = asString(params.stopReason);
     this.closeOpenTools(time);
     this.closeTurn({ stopReason, errored: stopReason === "error" }, time);
     // The sandbox never emits task_complete for successful runs (the terminal
@@ -232,12 +230,11 @@ export class RunTraceBuilder {
   }
 
   private startTool(update: Record<string, unknown>, time: Date): Context {
-    const toolCallId =
-      typeof update.toolCallId === "string" ? update.toolCallId : undefined;
+    const toolCallId = asString(update.toolCallId);
     const existing = toolCallId ? this.toolSpans.get(toolCallId) : undefined;
     if (existing) return existing.context;
 
-    const kind = typeof update.kind === "string" ? update.kind : "unknown";
+    const kind = asString(update.kind) ?? "unknown";
     const parentContext = this.turnContext ?? this.rootContext;
     const attributes: Attributes = { tool_kind: kind };
     if (toolCallId) attributes.tool_call_id = toolCallId;
@@ -263,8 +260,7 @@ export class RunTraceBuilder {
     update: Record<string, unknown>,
     time: Date,
   ): Context {
-    const toolCallId =
-      typeof update.toolCallId === "string" ? update.toolCallId : undefined;
+    const toolCallId = asString(update.toolCallId);
     const open = toolCallId ? this.toolSpans.get(toolCallId) : undefined;
     if (!open || !toolCallId) return this.currentContext();
 
@@ -283,8 +279,7 @@ export class RunTraceBuilder {
   private handleError(params: Record<string, unknown>, time: Date): Context {
     this.rootErrored = true;
     this.closeOpenTools(time, { interrupted: true });
-    const stopReason =
-      typeof params.stopReason === "string" ? params.stopReason : undefined;
+    const stopReason = asString(params.stopReason);
     this.closeTurn({ stopReason, errored: true }, time);
     // params.error is free text that can embed prompt or repo content, so
     // only the error's provenance is exported; the raw message stays in the

packages/core/src/sessions/sessionService.ts

@@ -5545,9 +5545,7 @@ export class SessionService {
             runState,
           );
         }
-        this.finalizeTerminalCloudTask(taskId, taskRunId, {
-          status: runStatus,
-        });
+        this.finalizeTerminalCloudTask(taskRunId, runStatus);
         this.stopCloudTaskWatch(taskId);
         return () => {};
       }
@@ -5744,9 +5742,7 @@ export class SessionService {
         runStatus,
         runState,
       );
-      this.finalizeTerminalCloudTask(taskId, taskRunId, {
-        status: runStatus,
-      });
+      this.finalizeTerminalCloudTask(taskRunId, runStatus);
       return () => {};
     }
 
@@ -6172,14 +6168,10 @@ export class SessionService {
         timestamp: Date.now(),
       });
     }
-    if (hasUserPrompt) {
-      // The real prompt has landed; the stash is no longer needed.
-      this.initialCloudOptimisticPrompt.delete(taskId);
-      this.d.store.clearTailOptimisticItems(taskRunId);
-    }
-    if (isTerminalRun) {
-      // A finished run gets no further echoes: any optimistic leftovers would
-      // otherwise linger as phantom tail items on the final transcript.
+    if (hasUserPrompt || isTerminalRun) {
+      // The stash is no longer needed once the real prompt lands - and a
+      // finished run gets no further echoes, so leftover optimistic items
+      // would otherwise linger as phantom tail items on the final transcript.
       this.initialCloudOptimisticPrompt.delete(taskId);
       this.d.store.clearTailOptimisticItems(taskRunId);
     }
@@ -6197,26 +6189,17 @@ export class SessionService {
     // Terminal hydration is different: it is the final transcript, so apply
     // it whenever the persisted chain has more lines than the local stream.
     const effectiveLineCount = Math.max(liveStreamLineCount, rawEntries.length);
-    if (isTerminalRun) {
-      if ((session.processedLineCount ?? 0) >= effectiveLineCount) {
-        this.surfacePersistedPendingPermissions(taskRunId, rawEntries);
-        this.pendingPermissionHydratedRuns.add(taskRunId);
-        return {
-          historyEntryCount: rawEntries.length,
-          liveStreamLineCount:
-            session.processedLineCount ?? liveStreamLineCount,
-        };
-      }
-    } else if (
-      session.processedLineCount !== undefined &&
-      session.processedLineCount > 0 &&
-      !isResumeRun
-    ) {
+    const alreadyApplied = isTerminalRun
+      ? (session.processedLineCount ?? 0) >= effectiveLineCount
+      : session.processedLineCount !== undefined &&
+        session.processedLineCount > 0 &&
+        !isResumeRun;
+    if (alreadyApplied) {
       this.surfacePersistedPendingPermissions(taskRunId, rawEntries);
       this.pendingPermissionHydratedRuns.add(taskRunId);
       return {
         historyEntryCount: rawEntries.length,
-        liveStreamLineCount: session.processedLineCount,
+        liveStreamLineCount: session.processedLineCount ?? liveStreamLineCount,
       };
     }
 
@@ -6247,7 +6230,7 @@ export class SessionService {
     // path otherwise sees delta <= 0 and never re-evaluates the tail.
     this.updatePromptStateFromEvents(taskRunId, events);
     if (isTerminalRun) {
-      this.clearTerminalCloudPromptState(taskId, taskRunId);
+      this.clearTerminalCloudPromptState(taskRunId);
     }
     return {
       historyEntryCount: rawEntries.length,
@@ -6256,28 +6239,18 @@ export class SessionService {
   }
 
   private finalizeTerminalCloudTask(
-    taskId: string,
     taskRunId: string,
-    fields: {
-      status?: TaskRunStatus;
-      stage?: string | null;
-      output?: Record<string, unknown> | null;
-      errorMessage?: string | null;
-      branch?: string | null;
-    },
+    status: TaskRunStatus | undefined,
   ): void {
-    this.d.store.updateCloudStatus(taskRunId, fields);
-    this.clearTerminalCloudPromptState(taskId, taskRunId);
+    this.d.store.updateCloudStatus(taskRunId, { status });
+    this.clearTerminalCloudPromptState(taskRunId);
   }
 
   /**
    * A terminal run can never flush its queue or answer a pending prompt;
    * leaving those set keeps the composer in a busy state forever.
    */
-  private clearTerminalCloudPromptState(
-    taskId: string,
-    taskRunId: string,
-  ): void {
+  private clearTerminalCloudPromptState(taskRunId: string): void {
     const session = this.d.store.getSessions()[taskRunId];
     if (
       !session ||
@@ -6286,7 +6259,7 @@ export class SessionService {
       return;
     }
 
-    this.d.store.clearMessageQueue(taskId);
+    this.d.store.clearMessageQueue(session.taskId);
     this.d.store.updateSession(taskRunId, {
       isPromptPending: false,
     });
@@ -7236,10 +7209,18 @@ export class SessionService {
       }
     }
 
-    if (update.kind === "snapshot" && !isTerminalStatus(update.status)) {
-      if (!this.isStaleNonTerminalCloudUpdate(taskRunId, update)) {
-        this.surfacePersistedPendingPermissions(taskRunId, update.newEntries);
-      }
+    // Evaluated once, before updateCloudStatus below can mutate cloudStatus.
+    const isStaleNonTerminalStatus = this.isStaleNonTerminalCloudUpdate(
+      taskRunId,
+      update,
+    );
+
+    if (
+      update.kind === "snapshot" &&
+      !isTerminalStatus(update.status) &&
+      !isStaleNonTerminalStatus
+    ) {
+      this.surfacePersistedPendingPermissions(taskRunId, update.newEntries);
     }
 
     // NOTE: Don't auto-flush on `!isPromptPending && queue.length > 0` here.
@@ -7252,7 +7233,7 @@ export class SessionService {
 
     // Update cloud status fields if present
     if (update.kind === "status" || update.kind === "snapshot") {
-      if (!this.isStaleNonTerminalCloudUpdate(taskRunId, update)) {
+      if (!isStaleNonTerminalStatus) {
         this.d.store.updateCloudStatus(taskRunId, {
           status: update.status,
           stage: update.stage,
@@ -7269,17 +7250,8 @@ export class SessionService {
       }
 
       if (isTerminalStatus(update.status)) {
-        // Clean up any pending resume messages that couldn't be sent
-        const session = this.d.store.getSessions()[taskRunId];
-        if (
-          session &&
-          (session.messageQueue.length > 0 || session.isPromptPending)
-        ) {
-          this.d.store.clearMessageQueue(session.taskId);
-          this.d.store.updateSession(taskRunId, {
-            isPromptPending: false,
-          });
-        }
+        // Pending resume messages can never be sent to a settled run.
+        this.clearTerminalCloudPromptState(taskRunId);
         this.stopCloudTaskWatch(update.taskId);
       }
     }

packages/core/src/sessions/sessionViewState.ts

@@ -4,6 +4,7 @@ import {
   type Task,
   type TaskRunStatus,
 } from "@posthog/shared/domain-types";
+import { resolveEffectiveCloudStatus } from "../task-detail/cloudRunState";
 
 export interface SessionViewState {
   isCloudRunNotTerminal: boolean;
@@ -27,15 +28,7 @@ export function deriveSessionViewState(
   workspace: Workspace | null,
   isCloud: boolean,
 ): SessionViewState {
-  const taskRunId = task.latest_run?.id;
-  const taskRunStatus = task.latest_run?.status ?? null;
-  const sessionMatchesLatestRun =
-    !!taskRunId && session?.taskRunId === taskRunId;
-  const cloudStatus = sessionMatchesLatestRun
-    ? isTerminalStatus(taskRunStatus)
-      ? taskRunStatus
-      : (session?.cloudStatus ?? taskRunStatus)
-    : (taskRunStatus ?? session?.cloudStatus ?? null);
+  const cloudStatus = resolveEffectiveCloudStatus(task, session);
   const isCloudRunTerminal = isCloud && isTerminalStatus(cloudStatus);
   const isCloudRunNotTerminal = isCloud && !isCloudRunTerminal;