feat: add traces to observability client and wire to runtime handlers

src/core/observability.test.ts

@@ -1,7 +1,6 @@
 import { describe, expect, test } from "bun:test";
 import {
   GetQueryResultsCommand,
-  ResourceNotFoundException,
   StartQueryCommand,
   type CloudWatchLogsClient,
 } from "@aws-sdk/client-cloudwatch-logs";
@@ -266,9 +265,12 @@ function insightsLogs(results: { field: string; value: string }[][]) {
   return { logs, queries };
 }
 
-describe("ObservabilityClient.listRuntimeTraces", () => {
-  const INPUT = {
-    runtimeId: "my_agent-AbC123XyZ9",
+const TRACE_SOURCE = {
+  logGroupName: "/aws/bedrock-agentcore/runtimes/my_agent-AbC123XyZ9-DEFAULT",
+};
+
+describe("ObservabilityClient.listTraces", () => {
+  const QUERY = {
     startTimeMs: 1_700_000_000_123,
     endTimeMs: 1_700_003_600_456,
     limit: 5,
@@ -277,12 +279,10 @@ describe("ObservabilityClient.listRuntimeTraces", () => {
   test("aggregates traces with a stats-by-traceId query over the runtime log group", async () => {
     const { logs, queries } = insightsLogs([]);
 
-    await clientWith(logs).listRuntimeTraces(INPUT, OPTIONS);
+    await clientWith(logs).listTraces(TRACE_SOURCE, QUERY, OPTIONS);
 
     expect(queries).toHaveLength(1);
-    expect(queries[0]!.logGroupNames).toEqual([
-      "/aws/bedrock-agentcore/runtimes/my_agent-AbC123XyZ9-DEFAULT",
-    ]);
+    expect(queries[0]!.logGroupNames).toEqual([TRACE_SOURCE.logGroupName]);
     // Epoch ms narrows to whole seconds.
     expect(queries[0]!.startTime).toBe(1_700_000_000);
     expect(queries[0]!.endTime).toBe(1_700_003_600);
@@ -311,7 +311,7 @@ describe("ObservabilityClient.listRuntimeTraces", () => {
       ],
     ]);
 
-    const traces = await clientWith(logs).listRuntimeTraces(INPUT, OPTIONS);
+    const traces = await clientWith(logs).listTraces(TRACE_SOURCE, QUERY, OPTIONS);
 
     expect(traces).toEqual([
       {
@@ -324,21 +324,10 @@ describe("ObservabilityClient.listRuntimeTraces", () => {
       { traceId: "def456", timestamp: "1700000002000", sessionId: undefined, spanCount: undefined },
     ]);
   });
-
-  test("translates a missing log group into invoked-yet guidance", async () => {
-    const logs = fakeLogs(async () => {
-      throw new ResourceNotFoundException({ message: "no such group", $metadata: {} });
-    });
-
-    await expect(clientWith(logs).listRuntimeTraces(INPUT, OPTIONS)).rejects.toThrow(
-      "Has the runtime been invoked yet?",
-    );
-  });
 });
 
-describe("ObservabilityClient.getRuntimeTrace", () => {
-  const INPUT = {
-    runtimeId: "my_agent-AbC123XyZ9",
+describe("ObservabilityClient.getTrace", () => {
+  const QUERY = {
     traceId: "68b2fabc0000000000abcdef",
     startTimeMs: 1_700_000_000_000,
     endTimeMs: 1_700_003_600_000,
@@ -350,7 +339,7 @@ describe("ObservabilityClient.getRuntimeTrace", () => {
     });
 
     await expect(
-      clientWith(logs).getRuntimeTrace({ ...INPUT, traceId: "not'a$trace" }, OPTIONS),
+      clientWith(logs).getTrace(TRACE_SOURCE, { ...QUERY, traceId: "not'a$trace" }, OPTIONS),
     ).rejects.toThrow("Invalid trace ID format. Expected a hex string (e.g., abc123def456).");
   });
 
@@ -367,7 +356,7 @@ describe("ObservabilityClient.getRuntimeTrace", () => {
       ],
     ]);
 
-    const records = await clientWith(logs).getRuntimeTrace(INPUT, OPTIONS);
+    const records = await clientWith(logs).getTrace(TRACE_SOURCE, QUERY, OPTIONS);
 
     expect(queries[0]!.queryString).toBe(
       "fields @timestamp, @message\n" +
@@ -388,8 +377,20 @@ describe("ObservabilityClient.getRuntimeTrace", () => {
   test("fails when the trace has no records", async () => {
     const { logs } = insightsLogs([]);
 
-    await expect(clientWith(logs).getRuntimeTrace(INPUT, OPTIONS)).rejects.toThrow(
+    await expect(clientWith(logs).getTrace(TRACE_SOURCE, QUERY, OPTIONS)).rejects.toThrow(
       "No trace data found for trace ID: 68b2fabc0000000000abcdef",
     );
   });
+
+  test("returns every record when the trace reaches the 10,000-record query limit", async () => {
+    const { logs } = insightsLogs(
+      Array.from({ length: 10_000 }, (_, index) => [
+        { field: "@message", value: `record-${index}` },
+      ]),
+    );
+
+    const records = await clientWith(logs).getTrace(TRACE_SOURCE, QUERY, OPTIONS);
+
+    expect(records).toHaveLength(10_000);
+  });
 });

src/core/observability.ts

@@ -1,6 +1,5 @@
 import {
   GetQueryResultsCommand,
-  ResourceNotFoundException,
   StartQueryCommand,
   type CloudWatchLogsClient,
   type ResultField,
@@ -17,22 +16,14 @@ import {
 } from "../errors";
 import type { ReadWriteJson } from "../io";
 import type { Project } from "../handlers/project/types";
-import type {
-  CoreObservabilityClient,
-  DeployedRuntime,
-  GetRuntimeTraceInput,
-  ListRuntimeTracesInput,
-  TraceRecord,
-  TraceSummary,
-} from "../handlers/runtime/types";
+import type { CoreObservabilityClient, DeployedRuntime } from "../handlers/runtime/types";
 import { AwsDeploymentTargetsSchema } from "../projectSchemas/aws-targets";
 import {
   CloudWatchClient,
   ObservabilityClient as GenericObservabilityClient,
 } from "./observability/index";
 import { isStackNotFound } from "./project/backends/cdk/environment";
-import type { AwsClients, CoreOptions } from "./types";
-import { toClientConfig } from "./utils";
+import type { AwsClients } from "./types";
 
 // Shared CloudWatch observability helpers. AgentCore Runtimes write their logs
 // and OTel telemetry to per-runtime CloudWatch log groups; both the eval flows
@@ -243,21 +234,18 @@ export interface ObservabilityClientDeps {
 }
 
 /**
- * Runtime-specific observability APIs retained for the existing trace and
- * project-resolution commands. Generic log reads are inherited from the new
- * shared observability client.
+ * Project-resolution APIs retained for the existing project command path.
+ * Shared observability operations are inherited from the generic client.
  */
 export class ObservabilityClient
   extends GenericObservabilityClient
   implements CoreObservabilityClient
 {
-  private readonly clients: AwsClients;
   private readonly readJson: ReadWriteJson;
   private readonly describeStackOutputs: DescribeStackOutputs;
 
   constructor(clients: AwsClients, deps: ObservabilityClientDeps) {
     super(new CloudWatchClient(clients));
-    this.clients = clients;
     this.readJson = deps.readJson;
     this.describeStackOutputs = deps.describeStackOutputs ?? describeStackOutputsWithSdk;
   }
@@ -320,124 +308,4 @@ export class ObservabilityClient
       targetName: target.name,
     };
   }
-
-  /**
-   * Lists the runtime's recent traces by aggregating its telemetry records with
-   * a Logs Insights `stats … by traceId` query (mirrors the old CLI's
-   * list-traces operation), newest first.
-   */
-  async listRuntimeTraces(
-    input: ListRuntimeTracesInput,
-    options: CoreOptions,
-  ): Promise<TraceSummary[]> {
-    const logGroupName = runtimeLogGroup(input.runtimeId, DEFAULT_ENDPOINT_QUALIFIER);
-    // Infrastructure records carry an empty traceId; excluding them before the
-    // aggregation keeps them from occupying one of the `limit` buckets (the old
-    // CLI filtered afterwards, silently returning one trace fewer).
-    const queryString =
-      `filter ispresent(traceId) and traceId != ""\n` +
-      `| stats earliest(@timestamp) as firstSeen, latest(@timestamp) as lastSeen, ` +
-      `count(*) as spanCount, earliest(attributes.session.id) as sessionId by traceId\n` +
-      `| sort lastSeen desc\n` +
-      `| limit ${Math.floor(input.limit)}`;
-
-    const rows = await this.runTraceQuery(input, logGroupName, queryString, options);
-
-    const traces: TraceSummary[] = [];
-    for (const row of rows) {
-      const fields = fieldMap(row);
-      if (!fields.traceId) continue;
-      traces.push({
-        traceId: fields.traceId,
-        timestamp: fields.lastSeen ?? fields.firstSeen ?? "unknown",
-        sessionId: fields.sessionId,
-        spanCount: fields.spanCount,
-      });
-    }
-    return traces;
-  }
-
-  /**
-   * Downloads every log record belonging to one trace, oldest first. The
-   * `@message` body is JSON-parsed when possible; other Insights fields pass
-   * through as returned.
-   */
-  async getRuntimeTrace(input: GetRuntimeTraceInput, options: CoreOptions): Promise<TraceRecord[]> {
-    if (!TRACE_ID_PATTERN.test(input.traceId)) {
-      throw new InputValidationError(
-        "Invalid trace ID format. Expected a hex string (e.g., abc123def456).",
-        { meta: { traceId: input.traceId } },
-      );
-    }
-
-    const logGroupName = runtimeLogGroup(input.runtimeId, DEFAULT_ENDPOINT_QUALIFIER);
-    const queryString =
-      `fields @timestamp, @message\n` +
-      `| filter traceId = '${sanitizeQueryValue(input.traceId)}'\n` +
-      `| sort @timestamp asc\n` +
-      `| limit 10000`;
-
-    const rows = await this.runTraceQuery(input, logGroupName, queryString, options);
-    if (rows.length === 0) {
-      throw new ResourceNotFoundError(`No trace data found for trace ID: ${input.traceId}`, {
-        meta: { traceId: input.traceId },
-      });
-    }
-
-    return rows.map((row) => {
-      const record: TraceRecord = fieldMap(row);
-      const message = record["@message"];
-      if (typeof message === "string") {
-        try {
-          record["@message"] = JSON.parse(message);
-        } catch {
-          // Keep the original string when the body is not valid JSON.
-        }
-      }
-      return record;
-    });
-  }
-
-  private async runTraceQuery(
-    input: { runtimeId: string; startTimeMs: number; endTimeMs: number },
-    logGroupName: string,
-    queryString: string,
-    options: CoreOptions,
-  ): Promise<ResultField[][]> {
-    const logs = this.clients.logs(toClientConfig(options));
-    const startSec = Math.floor(input.startTimeMs / 1000);
-    const endSec = Math.floor(input.endTimeMs / 1000);
-    try {
-      return await runInsightsQuery(logs, [logGroupName], queryString, startSec, endSec);
-    } catch (error) {
-      if (error instanceof ResourceNotFoundException) {
-        throw missingLogGroupError(input.runtimeId, logGroupName, error);
-      }
-      throw error;
-    }
-  }
-}
-
-// Trace ids are hex strings, optionally dash-separated (mirrors the old CLI).
-const TRACE_ID_PATTERN = /^[a-fA-F0-9-]+$/;
-
-// fieldMap flattens one Insights result row into a name -> value record.
-function fieldMap(row: ResultField[]): Record<string, string> {
-  const fields: Record<string, string> = {};
-  for (const field of row) {
-    if (field.field && field.value !== undefined) fields[field.field] = field.value;
-  }
-  return fields;
-}
-
-function missingLogGroupError(
-  runtimeId: string,
-  logGroupName: string,
-  cause?: unknown,
-): ResourceNotFoundError {
-  return new ResourceNotFoundError(
-    `No logs found for runtime '${runtimeId}': log group ${logGroupName} does not exist. ` +
-      `Has the runtime been invoked yet?`,
-    { cause, meta: { runtimeId, logGroupName } },
-  );
 }

src/core/observability/client.ts

@@ -2,12 +2,22 @@ import type { CoreOptions } from "../types";
 import { CloudWatchClient } from "./cloudWatchClient";
 import type {
   CloudWatchLogEvent,
+  GetTraceQuery,
   InsightsQuery,
   InsightsQueryRow,
+  ListTracesQuery,
   LogSearchQuery,
   LogSource,
   LogTailQuery,
+  TraceRecord,
+  TraceSummary,
 } from "./types";
+import {
+  getTraceInsightsQuery,
+  listTracesInsightsQuery,
+  normalizeTraceRecords,
+  normalizeTraceSummaries,
+} from "./traces";
 
 /** Shared observability API over explicit CloudWatch log-group targets. */
 export class ObservabilityClient {
@@ -39,4 +49,24 @@ export class ObservabilityClient {
   ): Promise<InsightsQueryRow[]> {
     return this.cloudWatch.queryLogs(source, query, options, signal);
   }
+
+  async listTraces(
+    source: LogSource,
+    query: ListTracesQuery,
+    options: CoreOptions,
+    signal?: AbortSignal,
+  ): Promise<TraceSummary[]> {
+    const rows = await this.queryLogs(source, listTracesInsightsQuery(query), options, signal);
+    return normalizeTraceSummaries(rows);
+  }
+
+  async getTrace(
+    source: LogSource,
+    query: GetTraceQuery,
+    options: CoreOptions,
+    signal?: AbortSignal,
+  ): Promise<TraceRecord[]> {
+    const rows = await this.queryLogs(source, getTraceInsightsQuery(query), options, signal);
+    return normalizeTraceRecords(rows, query.traceId);
+  }
 }

src/core/observability/index.ts

@@ -6,11 +6,16 @@ export {
   sanitizeQueryValue,
   type InsightsRowLimit,
 } from "./insights";
+export { TRACE_RECORD_LIMIT } from "./traces";
 export type {
   CloudWatchLogEvent,
+  GetTraceQuery,
   InsightsQuery,
   InsightsQueryRow,
+  ListTracesQuery,
   LogSearchQuery,
   LogSource,
   LogTailQuery,
+  TraceRecord,
+  TraceSummary,
 } from "./types";