clean up invocation message types and improve callback typing

.gitignore

@@ -21,6 +21,7 @@ node_modules/
 Thumbs.db
 
 # IDE / editor
+.vscode/
 .idea/
 *.swp
 *.swo

README.md

@@ -32,7 +32,7 @@ async def main():
     invocation = InvocationMessage(
         message_type="invocation_message",
         source_id="my-run-001",
-        preferred_model="gpt-4.1",
+        model="gpt-4.1",
         system_prompt="You are a helpful assistant.",
         user_prompt="What is 2 + 2?",
         output_format_json_schema_str=json.dumps({
@@ -43,7 +43,6 @@ async def main():
             },
         }),
         secrets_to_redact=[],
-        agent_env={},
     )
 
     async def on_result(result: ResultMessageInterface):

docs/content/docs/quickstart.mdx

@@ -41,12 +41,9 @@ invoker = RuntimeUseInvoker(client)
 # Build an invocation
 invocation = InvocationMessage(
     message_type="invocation_message",
-    source_id="my-source",
     system_prompt="You are a helpful assistant.",
     user_prompt="Run the tests.",
-    output_format_json_schema_str='{"type":"json_schema","schema":{"type":"object"}}',
-    secrets_to_redact=[],
-    agent_env={},
+    output_format_json_schema_str='{"type":"json_schema","schema":{"type":"object"}}'
 )
 
 # Invoke and handle the result

packages/runtimeuse-client-python/README.md

@@ -30,11 +30,11 @@ async def main():
     invocation = InvocationMessage(
         message_type="invocation_message",
         source_id="my-run-001",
+        model="gpt-4.1",
         system_prompt="You are a helpful assistant.",
         user_prompt="Do the thing and return the result.",
         output_format_json_schema_str='{"type":"json_schema","schema":{"type":"object"}}',
         secrets_to_redact=["sk-secret-key"],
-        agent_env={"API_KEY": "sk-secret-key"},
     )
 
     async def on_result(result: ResultMessageInterface):
@@ -70,7 +70,7 @@ await client.invoke(
     on_result_message=on_result,
     result_message_cls=ResultMessageInterface,
     on_assistant_message=on_assistant,       # optional
-    on_artifact_upload_request=on_artifact,  # optional -- return (presigned_url, content_type)
+    on_artifact_upload_request=on_artifact,  # optional -- return ArtifactUploadResult
     on_error_message=on_error,               # optional
     is_cancelled=check_cancelled,            # optional -- async () -> bool
     timeout=300,                             # optional -- seconds
@@ -82,10 +82,12 @@ await client.invoke(
 When the agent runtime requests an artifact upload, provide a callback that returns a presigned URL and content type. The client sends the response back automatically.
 
 ```python
-async def on_artifact(request: ArtifactUploadRequestMessageInterface) -> tuple[str, str]:
+from runtimeuse import ArtifactUploadResult
+
+async def on_artifact(request: ArtifactUploadRequestMessageInterface) -> ArtifactUploadResult:
     presigned_url = await my_storage.create_presigned_url(request.filename)
     content_type = guess_content_type(request.filename)
-    return presigned_url, content_type
+    return ArtifactUploadResult(presigned_url=presigned_url, content_type=content_type)
 ```
 
 ### Cancellation

packages/runtimeuse-client-python/examples/runtime-client-example.py

@@ -22,7 +22,7 @@ async def main():
     invocation = InvocationMessage(
         message_type="invocation_message",
         source_id="my-source",
-        preferred_model="gpt-5.4",
+        model="gpt-5.4",
         pre_agent_invocation_commands=[
             CommandInterface(
                 command="echo 'Hello, world!'",
@@ -36,7 +36,6 @@ async def main():
             {"type": "json_schema", "schema": Answer.model_json_schema()}
         ),
         secrets_to_redact=[],
-        agent_env={},
     )
 
     async def on_result(result: ResultMessageInterface):

packages/runtimeuse-client-python/pyproject.toml

@@ -20,5 +20,12 @@ Homepage = "https://github.com/getlark/runtimeuse"
 Repository = "https://github.com/getlark/runtimeuse"
 Issues = "https://github.com/getlark/runtimeuse/issues"
 
+[project.optional-dependencies]
+dev = [
+    "daytona",
+    "e2b",
+    "e2b-code-interpreter",
+]
+
 [tool.hatch.build.targets.wheel]
 packages = ["src/runtimeuse_client"]

packages/runtimeuse-client-python/src/runtimeuse_client/__init__.py

@@ -12,6 +12,11 @@
     ArtifactUploadResponseMessageInterface,
     ErrorMessageInterface,
     CancelMessage,
+    ArtifactUploadResult,
+    OnAssistantMessageCallback,
+    OnArtifactUploadRequestCallback,
+    OnErrorMessageCallback,
+    IsCancelledCallback,
 )
 
 __all__ = [
@@ -29,4 +34,9 @@
     "ArtifactUploadResponseMessageInterface",
     "ErrorMessageInterface",
     "CancelMessage",
+    "ArtifactUploadResult",
+    "OnAssistantMessageCallback",
+    "OnArtifactUploadRequestCallback",
+    "OnErrorMessageCallback",
+    "IsCancelledCallback",
 ]

packages/runtimeuse-client-python/src/runtimeuse_client/client.py

@@ -14,6 +14,10 @@
     AssistantMessageInterface,
     ArtifactUploadRequestMessageInterface,
     ArtifactUploadResponseMessageInterface,
+    OnAssistantMessageCallback,
+    OnArtifactUploadRequestCallback,
+    OnErrorMessageCallback,
+    IsCancelledCallback,
 )
 from .exceptions import CancelledException
 
@@ -50,22 +54,13 @@ def __init__(
     async def invoke(
         self,
         invocation: InvocationMessage,
+        # this should be response instead?
         on_result_message: Callable[[T], Awaitable[None]],
         result_message_cls: Type[T],
-        on_assistant_message: (
-            Callable[[AssistantMessageInterface], Awaitable[None]] | None
-        ) = None,
-        on_artifact_upload_request: (
-            Callable[
-                [ArtifactUploadRequestMessageInterface],
-                Awaitable[tuple[str, str]],
-            ]
-            | None
-        ) = None,
-        on_error_message: (
-            Callable[[ErrorMessageInterface], Awaitable[None]] | None
-        ) = None,
-        is_cancelled: Callable[[], Awaitable[bool]] | None = None,
+        on_assistant_message: OnAssistantMessageCallback | None = None,
+        on_artifact_upload_request: OnArtifactUploadRequestCallback | None = None,
+        on_error_message: OnErrorMessageCallback | None = None,
+        is_cancelled: IsCancelledCallback | None = None,
         timeout: float | None = None,
         logger: logging.Logger | None = None,
     ) -> None:
@@ -79,8 +74,8 @@ async def invoke(
             on_assistant_message: Optional async callback invoked when an assistant_message
                 is received.
             on_artifact_upload_request: Optional async callback invoked when an
-                artifact_upload_request_message is received. Should return a
-                (presigned_url, content_type) tuple; the client will send the
+                artifact_upload_request_message is received. Should return an
+                ArtifactUploadResult; the client will send the
                 artifact_upload_response_message back to the agent runtime automatically.
             on_error_message: Optional async callback invoked when an error_message is
                 received.
@@ -164,15 +159,15 @@ async def invoke(
                                 message
                             )
                         )
-                        presigned_url, content_type = await on_artifact_upload_request(
+                        upload_result = await on_artifact_upload_request(
                             artifact_upload_request_message_interface
                         )
                         artifact_upload_response_message_interface = ArtifactUploadResponseMessageInterface(
                             message_type="artifact_upload_response_message",
                             filename=artifact_upload_request_message_interface.filename,
                             filepath=artifact_upload_request_message_interface.filepath,
-                            presigned_url=presigned_url,
-                            content_type=content_type,
+                            presigned_url=upload_result.presigned_url,
+                            content_type=upload_result.content_type,
                         )
                         await send_queue.put(
                             artifact_upload_response_message_interface.model_dump(

packages/runtimeuse-client-python/src/runtimeuse_client/types.py

@@ -1,4 +1,4 @@
-from typing import Any, Literal
+from typing import Any, Awaitable, Callable, Literal
 
 from pydantic import BaseModel
 from pydantic.fields import Field
@@ -28,17 +28,17 @@ class CommandInterface(BaseModel):
 
 class InvocationMessage(BaseModel):
     message_type: Literal["invocation_message"]
-    source_id: str
+    source_id: str | None = None
     system_prompt: str
     user_prompt: str
     output_format_json_schema_str: str
-    secrets_to_redact: list[str]
-    agent_env: dict[str, str]
+    secrets_to_redact: list[str] = Field(default_factory=list)
     artifacts_dir: str | None = None
     pre_agent_invocation_commands: list[CommandInterface] | None = None
     post_agent_invocation_commands: list[CommandInterface] | None = None
-    preferred_model: str
-    runtime_environment_downloadables: (
+    model: str
+
+    pre_agent_downloadables: (
         list[RuntimeEnvironmentDownloadableInterface] | None
     ) = None
 
@@ -76,3 +76,16 @@ class ErrorMessageInterface(AgentRuntimeMessageInterface):
 
 class CancelMessage(BaseModel):
     message_type: Literal["cancel_message"]
+
+
+class ArtifactUploadResult(BaseModel):
+    presigned_url: str
+    content_type: str
+
+
+OnAssistantMessageCallback = Callable[[AssistantMessageInterface], Awaitable[None]]
+OnArtifactUploadRequestCallback = Callable[
+    [ArtifactUploadRequestMessageInterface], Awaitable[ArtifactUploadResult]
+]
+OnErrorMessageCallback = Callable[[ErrorMessageInterface], Awaitable[None]]
+IsCancelledCallback = Callable[[], Awaitable[bool]]

packages/runtimeuse-client-python/test/conftest.py

@@ -43,11 +43,10 @@ def _make_invocation(**overrides: Any) -> InvocationMessage:
         message_type="invocation_message",
         source_id="test-001",
         system_prompt="You are a good assistant.",
-        preferred_model="gpt-5.4`",
         user_prompt="Do something.",
         output_format_json_schema_str='{"type":"object"}',
         secrets_to_redact=[],
-        agent_env={},
+        model="gpt-4o",
     )
     defaults.update(overrides)
     return InvocationMessage.model_validate(defaults)

packages/runtimeuse-client-python/test/test_client.py

@@ -8,6 +8,7 @@
     ResultMessageInterface,
     AssistantMessageInterface,
     ArtifactUploadRequestMessageInterface,
+    ArtifactUploadResult,
     ErrorMessageInterface,
     CancelledException,
 )
@@ -179,9 +180,12 @@ async def test_artifact_upload_handshake(self, fake_transport, invocation):
 
         async def on_artifact(
             req: ArtifactUploadRequestMessageInterface,
-        ) -> tuple[str, str]:
+        ) -> ArtifactUploadResult:
             assert req.filename == "screenshot.png"
-            return "https://s3.example.com/presigned", "image/png"
+            return ArtifactUploadResult(
+                presigned_url="https://s3.example.com/presigned",
+                content_type="image/png",
+            )
 
         await client.invoke(
             invocation=invocation,

packages/runtimeuse/README.md

@@ -159,7 +159,7 @@ wss.on("connection", (ws) => {
 
 When a client sends an `invocation_message`, the session:
 
-1. **Downloads runtime files** -- if `runtime_environment_downloadables` is set, fetches and extracts them
+1. **Downloads runtime files** -- if `pre_agent_downloadables` is set, fetches and extracts them
 2. **Runs pre-commands** -- if `pre_agent_invocation_commands` is set, executes them. If it exits 0, execution continues to the next command or the agent. Any other non-zero exit code sends an error message and terminates the invocation.
 3. **Calls `handler.run()`** -- your agent logic runs with the invocation context and a `MessageSender`
 4. **Sends `result_message`** -- the `AgentResult` from your handler is sent back to the client

packages/runtimeuse/src/agent-handler.ts

@@ -6,7 +6,6 @@ export interface AgentInvocation {
   outputFormat: { type: "json_schema"; schema: Record<string, unknown> };
   model: string;
   secrets: string[];
-  env: Record<string, string>;
   signal: AbortSignal;
   logger: Logger;
 }

packages/runtimeuse/src/claude-handler.ts

@@ -40,8 +40,6 @@ export const claudeHandler: AgentHandler = {
           model: invocation.model,
           outputFormat: invocation.outputFormat,
           abortController,
-          cwd: process.cwd(),
-          env: { ...process.env, ...invocation.env },
           tools: { type: "preset", preset: "claude_code" },
           permissionMode: "bypassPermissions",
           allowDangerouslySkipPermissions: true,

packages/runtimeuse/src/invocation-runner.test.ts

@@ -61,7 +61,7 @@ const BASE_INVOCATION_MESSAGE: InvocationMessage = {
     type: "json_schema",
     schema: { type: "object", properties: { ok: { type: "boolean" } } },
   }),
-  preferred_model: "test-model",
+  model: "test-model",
 };
 
 function createRunner(overrides?: Partial<InvocationMessage>) {
@@ -113,9 +113,8 @@ describe("InvocationRunner", () => {
           type: "json_schema",
           schema: { type: "object", properties: { ok: { type: "boolean" } } },
         },
-        model: message.preferred_model,
+        model: message.model,
         secrets: message.secrets_to_redact,
-        env: {},
         signal: abortController.signal,
         logger,
       }),
@@ -147,7 +146,7 @@ describe("InvocationRunner", () => {
     });
 
     const { runner, message } = createRunner({
-      runtime_environment_downloadables: [
+      pre_agent_downloadables: [
         {
           download_url: "https://example.com/runtime.tar.gz",
           working_dir: "/tmp",

packages/runtimeuse/src/invocation-runner.ts

@@ -37,9 +37,8 @@ export class InvocationRunner {
         systemPrompt: message.system_prompt,
         userPrompt: message.user_prompt,
         outputFormat,
-        model: message.preferred_model,
+        model: message.model,
         secrets: message.secrets_to_redact,
-        env: message.agent_env ?? {},
         signal: abortController.signal,
         logger,
       },
@@ -70,10 +69,10 @@ export class InvocationRunner {
   private async downloadRuntimeEnvironment(
     message: InvocationMessage,
   ): Promise<void> {
-    if (!message.runtime_environment_downloadables) return;
+    if (!message.pre_agent_downloadables) return;
 
     this.config.logger.log("Downloading runtime environment downloadables...");
-    for (const downloadable of message.runtime_environment_downloadables) {
+    for (const downloadable of message.pre_agent_downloadables) {
       await this.downloadHandler.download(
         downloadable.download_url,
         downloadable.working_dir,

packages/runtimeuse/src/openai-handler.ts

@@ -19,8 +19,6 @@ export const openaiHandler: AgentHandler = {
     invocation: AgentInvocation,
     sender: MessageSender,
   ): Promise<AgentResult> {
-    Object.assign(process.env, invocation.env);
-
     const strictSchema = ensureStrictSchema(invocation.outputFormat.schema);
     const outputType = zod.fromJSONSchema(strictSchema) as AgentOutputType;
 

packages/runtimeuse/src/session.test.ts

@@ -104,7 +104,7 @@ const INVOCATION_MSG = {
     type: "json_schema",
     schema: { type: "object" },
   }),
-  preferred_model: "test-model",
+  model: "test-model",
 };
 
 describe("WebSocketSession", () => {
@@ -310,7 +310,7 @@ describe("WebSocketSession", () => {
       const done = session.run();
       sendMessage(ws, {
         ...INVOCATION_MSG,
-        runtime_environment_downloadables: [
+        pre_agent_downloadables: [
           {
             download_url: "https://example.com/test.zip",
             working_dir: "/tmp/test",