From bef999d3fae14094319bd62bd31d7004ecf610da Mon Sep 17 00:00:00 2001
From: Jack Kleeman <jackkleeman@gmail.com>
Date: Fri, 20 Dec 2024 10:41:06 +0000
Subject: [PATCH] Include request struct in logger context

So you can get access to headers, extra args
---
 packages/restate-sdk/src/context_impl.ts      | 24 +-----------------
 .../src/endpoint/handlers/generic.ts          | 25 +++++++++++++++++--
 .../src/logging/logger_transport.ts           |  3 +++
 3 files changed, 27 insertions(+), 25 deletions(-)

diff --git a/packages/restate-sdk/src/context_impl.ts b/packages/restate-sdk/src/context_impl.ts
index 2e56c25e..5c843142 100644
--- a/packages/restate-sdk/src/context_impl.ts
+++ b/packages/restate-sdk/src/context_impl.ts
@@ -53,7 +53,6 @@ import type {
 } from "@restatedev/restate-sdk-core";
 import { serde } from "@restatedev/restate-sdk-core";
 import { RandImpl } from "./utils/rand.js";
-import type { Headers } from "./endpoint/handlers/generic.js";
 import type {
   ReadableStreamDefaultReader,
   WritableStreamDefaultWriter,
@@ -67,7 +66,6 @@ export type InternalCombineablePromise<T> = CombineablePromise<T> & {
 };
 
 export class ContextImpl implements ObjectContext, WorkflowContext {
-  private readonly invocationRequest: Request;
   public readonly rand: Rand;
 
   public readonly date: ContextDate = {
@@ -86,31 +84,11 @@ export class ContextImpl implements ObjectContext, WorkflowContext {
     readonly input: vm.WasmInput,
     public readonly console: Console,
     public readonly handlerKind: HandlerKind,
-    attemptHeaders: Headers,
-    extraArgs: unknown[],
+    private readonly invocationRequest: Request,
     private readonly invocationEndPromise: CompletablePromise<void>,
     private readonly inputReader: ReadableStreamDefaultReader<Uint8Array>,
     private readonly outputWriter: WritableStreamDefaultWriter<Uint8Array>
   ) {
-    this.invocationRequest = {
-      id: input.invocation_id,
-      headers: input.headers.reduce((headers, { key, value }) => {
-        headers.set(key, value);
-        return headers;
-      }, new Map()),
-      attemptHeaders: Object.entries(attemptHeaders).reduce(
-        (headers, [key, value]) => {
-          if (value !== undefined) {
-            headers.set(key, value instanceof Array ? value[0] : value);
-          }
-          return headers;
-        },
-        new Map()
-      ),
-      body: input.input,
-      extraArgs,
-    };
-
     this.rand = new RandImpl(input.invocation_id, () => {
       if (coreVm.is_inside_run()) {
         throw new Error(
diff --git a/packages/restate-sdk/src/endpoint/handlers/generic.ts b/packages/restate-sdk/src/endpoint/handlers/generic.ts
index b2671777..dd64967e 100644
--- a/packages/restate-sdk/src/endpoint/handlers/generic.ts
+++ b/packages/restate-sdk/src/endpoint/handlers/generic.ts
@@ -23,6 +23,7 @@ import type { EndpointBuilder } from "../endpoint_builder.js";
 import { type ReadableStream, TransformStream } from "node:stream/web";
 import { OnceStream } from "../../utils/streams.js";
 import { ContextImpl } from "../../context_impl.js";
+import type { Request } from "../../context.js";
 import * as vm from "./vm/sdk_shared_core_wasm_bindings.js";
 import { CompletablePromise } from "../../utils/completable_promise.js";
 import { HandlerKind } from "../../types/rpc.js";
@@ -252,12 +253,33 @@ export class GenericHandler implements RestateHandler {
 
     // Get input
     const input = coreVm.sys_input();
+
+    const invocationRequest: Request = {
+      id: input.invocation_id,
+      headers: input.headers.reduce((headers, { key, value }) => {
+        headers.set(key, value);
+        return headers;
+      }, new Map()),
+      attemptHeaders: Object.entries(headers).reduce(
+        (headers, [key, value]) => {
+          if (value !== undefined) {
+            headers.set(key, value instanceof Array ? value[0] : value);
+          }
+          return headers;
+        },
+        new Map()
+      ),
+      body: input.input,
+      extraArgs,
+    };
+
     // Prepare logger
     const loggerContext = new LoggerContext(
       input.invocation_id,
       handler.component().name(),
       handler.name(),
       handler.kind() === HandlerKind.SERVICE ? undefined : input.key,
+      invocationRequest,
       additionalContext
     );
     const ctxLogger = createLogger(
@@ -294,8 +316,7 @@ export class GenericHandler implements RestateHandler {
       input,
       ctxLogger,
       handler.kind(),
-      headers,
-      extraArgs,
+      invocationRequest,
       invocationEndPromise,
       inputReader,
       outputWriter
diff --git a/packages/restate-sdk/src/logging/logger_transport.ts b/packages/restate-sdk/src/logging/logger_transport.ts
index f1e5c9d2..a31e77cf 100644
--- a/packages/restate-sdk/src/logging/logger_transport.ts
+++ b/packages/restate-sdk/src/logging/logger_transport.ts
@@ -11,6 +11,8 @@
 
 /* eslint-disable @typescript-eslint/no-explicit-any */
 
+import type { Request } from "../context.js";
+
 /**
  * Logger level.
  */
@@ -73,6 +75,7 @@ export class LoggerContext {
     readonly serviceName: string,
     readonly handlerName: string,
     readonly key?: string,
+    readonly request?: Request,
     readonly additionalContext?: { [name: string]: string }
   ) {
     this.invocationTarget =