Clarify CodeMode final-expression returns (#370)

* Prevent empty CodeMode results through clearer guidance

* Make CodeMode return guidance unambiguous

docs/code-mode.md

@@ -153,13 +153,23 @@ That fold-in grows `run_code`'s description, which invalidates the prompt-cache
 
 ## Return values
 
-The last expression in the snippet is automatically captured as the return value -- the model does not need to `print()`. Reserve `print()` for supplementary logging: printed text is surfaced separately, wrapped alongside the last-expression result.
+The last expression in the snippet is automatically captured as the return value -- the model does not need to `print()`. An assignment stores a value in the REPL but does not return it. A final expression that evaluates to `None` is also treated as no result. Without a non-`None` final expression or print output, `run_code` returns `{}`. Put the assigned name on the final line:
+
+```python
+result = await get_weather(city='Paris')
+result
+```
+
+Reserve `print()` for supplementary logging: printed text is surfaced separately, wrapped alongside the last-expression result.
 
 | Scenario | Return |
 |---|---|
-| No print output | Last expression value |
-| With print output | `{'output': '<printed text>', 'result': <last expression>}` |
-| Multimodal content (e.g. images) | Returned natively for model processing |
+| Non-`None` final expression with no print output | Last expression value |
+| Final assignment or `None` result with no print output | `{}` |
+| Print output with no final expression or a `None` result | `{'output': '<printed text>'}` |
+| Print output with a plain, non-`None` final expression | `{'output': '<printed text>', 'result': <last expression>}` |
+| Multimodal final expression with no print output | Returned natively for model processing |
+| Print output with a multimodal final expression | List with printed text followed by native multimodal content |
 
 ## REPL state
 

pydantic_ai_harness/code_mode/README.md

@@ -44,6 +44,8 @@ print(result.output)
 The model writes code like:
 
 ```python
+import asyncio
+
 paris, tokyo = await asyncio.gather(
     get_weather(city='Paris'),
     get_weather(city='Tokyo'),
@@ -144,13 +146,21 @@ become callable functions inside `run_code`. Tools without
 
 ## Return values
 
-The last expression in the code snippet is automatically captured as the return value -- the model does not need to `print()`.
+The last expression in the code snippet is automatically captured as the return value -- the model does not need to `print()`. An assignment stores a value in the REPL but does not return it. A final expression that evaluates to `None` is also treated as no result. Without a non-`None` final expression or print output, `run_code` returns `{}`. Put the assigned name on the final line:
+
+```python
+result = await get_weather(city='Paris')
+result
+```
 
 | Scenario | Return |
 |---|---|
-| No print output | Last expression value |
-| With print output | `{"output": "<printed text>", "result": <last expression>}` |
-| Multimodal content (e.g. images) | Returned natively for model processing |
+| Non-`None` final expression with no print output | Last expression value |
+| Final assignment or `None` result with no print output | `{}` |
+| Print output with no final expression or a `None` result | `{"output": "<printed text>"}` |
+| Print output with a plain, non-`None` final expression | `{"output": "<printed text>", "result": <last expression>}` |
+| Multimodal final expression with no print output | Returned natively for model processing |
+| Print output with a multimodal final expression | List with printed text followed by native multimodal content |
 
 ## REPL state
 

pydantic_ai_harness/code_mode/_toolset.py

@@ -85,7 +85,7 @@ Write and run Python code in a sandboxed environment.
 The sandbox uses Monty, a subset of Python. Key restrictions:
 - **No classes**: class definitions are not supported
 - **No third-party libraries**: only the standard library modules listed below can be used
-- **Importable standard library modules**: `sys`, `typing`, `asyncio`, `math`, `json`, `re`, `datetime`, `os`, `pathlib`. These must be imported before use, just like in regular Python. For example: `import asyncio` then `results = await asyncio.gather(tool_one(...), tool_two(...))`."""
+- **Importable standard library modules**: `sys`, `typing`, `asyncio`, `math`, `json`, `re`, `datetime`, `os`, `pathlib`. These must be imported before use, just like in regular Python. For example: `import asyncio` then `await asyncio.gather(tool_one(...), tool_two(...))`."""
 
 # Timing/OS restriction line, swapped depending on what host access the agent
 # configured. Three states, because `mount` and `os` enable different things:
@@ -116,11 +116,23 @@ _RUN_CODE_DESCRIPTION_TAIL = """\
 State is preserved between calls (REPL-style). Set `restart: true` to reset state.
 
 The last expression's value is automatically captured as the return value -- you do **not** need to \
-`print()` it. Avoid `print()` for return values as it produces Python string representations, not \
-structured data. Use `print()` only for supplementary logging or debug output.
+`print()` it. End the snippet with the value to return as a bare expression. A final assignment stores \
+the value but does not return it. A final expression that evaluates to `None` is treated as no result. \
+Without a non-`None` final expression or print output, `run_code` returns `{}`. For example:
 
-Returns the last expression's value directly. If `print()` was also called, returns \
-`{"output": "<printed text>", "result": <last expression>}`.\
+```python
+result = some_expression
+result
+```
+
+Avoid `print()` for return values as it produces Python string representations, not structured data. \
+Use `print()` only for supplementary logging or debug output.
+
+Returns a non-`None` last expression's value directly when nothing is printed. With `print()` output \
+and no non-`None` final expression, returns `{"output": "<printed text>"}`. With `print()` output and a \
+plain, non-`None` final expression, returns \
+`{"output": "<printed text>", "result": <last expression>}`. With `print()` output and a multimodal \
+final expression, returns a list with the printed text followed by the native content.\
 """
 
 
@@ -149,15 +161,15 @@ def _functions_header(*, has_sync: bool, has_async: bool) -> str:
     if has_async and not has_sync:
         return base + (
             ' All tool functions are async: invoke them with `await`,'
-            ' e.g. `result = await tool_name(arg=value)`.'
+            ' e.g. `await tool_name(arg=value)`.'
             ' Calling without `await` returns an unresolved future, not the value.'
         )
     if has_sync and not has_async:
-        return base + (' All tool functions are synchronous: call them directly, e.g. `result = tool_name(arg=value)`.')
+        return base + (' All tool functions are synchronous: call them directly, e.g. `tool_name(arg=value)`.')
     return base + (
         ' Async functions (`async def`) must be invoked with `await`,'
-        ' e.g. `result = await tool_name(arg=value)`.'
-        ' Sync functions (`def`) are called directly, e.g. `result = tool_name(arg=value)`.'
+        ' e.g. `await tool_name(arg=value)`.'
+        ' Sync functions (`def`) are called directly, e.g. `tool_name(arg=value)`.'
     )
 
 

tests/code_mode/test_code_mode.py

@@ -236,6 +236,48 @@ class TestCodeMode:
         # The base description must tell the model to await tool calls.
         assert 'await' in description
 
+    async def test_run_code_description_explains_final_expression_return(self) -> None:
+        """The model is told how to return a value after assigning it."""
+        wrapper = CodeMode[object]().get_wrapper_toolset(_build_function_toolset(add))
+        assert isinstance(wrapper, CodeModeToolset)
+
+        tools = await wrapper.get_tools(build_run_context(None))
+        description = tools['run_code'].tool_def.description
+
+        assert description is not None
+        assert 'End the snippet with the value to return as a bare expression.' in description
+        assert 'result = some_expression\nresult' in description
+        assert 'Without a non-`None` final expression or print output, `run_code` returns `{}`.' in description
+        assert 'A final expression that evaluates to `None` is treated as no result.' in description
+        assert 'results = await asyncio.gather' not in description
+        assert 'With `print()` output and no non-`None` final expression' in description
+        assert 'With `print()` output and a plain, non-`None` final expression' in description
+        assert 'With `print()` output and a multimodal final expression' in description
+
+    async def test_run_code_function_examples_are_expressions(self) -> None:
+        """Async, sync, and mixed function examples do not end on assignments."""
+        cases: list[tuple[FunctionToolset[object], tuple[str, ...]]] = [
+            (_build_function_toolset(add), ('e.g. `await tool_name(arg=value)`.',)),
+            (
+                FunctionToolset[object](tools=[Tool(add, sequential=True)]),
+                ('e.g. `tool_name(arg=value)`.',),
+            ),
+            (
+                FunctionToolset[object](tools=[Tool(add, sequential=True), Tool(greet)]),
+                ('e.g. `await tool_name(arg=value)`.', 'e.g. `tool_name(arg=value)`.'),
+            ),
+        ]
+
+        for toolset, expected_examples in cases:
+            wrapper = CodeMode[object]().get_wrapper_toolset(toolset)
+            assert isinstance(wrapper, CodeModeToolset)
+
+            description = (await wrapper.get_tools(build_run_context(None)))['run_code'].tool_def.description
+
+            assert description is not None
+            assert all(example in description for example in expected_examples)
+            assert 'e.g. `result =' not in description
+
     async def test_run_code_executes_call_through_monty(self) -> None:
         """End-to-end: `run_code` runs Python in Monty and dispatches to a sync wrapped tool."""
         toolset = _build_function_toolset(add)
@@ -390,6 +432,19 @@ class TestCodeMode:
         # No print output → result returned directly (not wrapped in a dict).
         assert result.return_value == 3
 
+    async def test_run_code_treats_none_as_no_expression_result(self) -> None:
+        """A final `None` uses the same return shapes as no final expression."""
+        wrapper = CodeMode[object]().get_wrapper_toolset(_build_function_toolset(add))
+        assert isinstance(wrapper, CodeModeToolset)
+        ctx = await build_ctx(None, wrapper)
+        tools = await wrapper.get_tools(ctx)
+
+        result = await wrapper.call_tool('run_code', {'code': 'None'}, ctx, tools['run_code'])
+        assert result.return_value == {}
+
+        printed = await wrapper.call_tool('run_code', {'code': 'print("done")\nNone'}, ctx, tools['run_code'])
+        assert printed.return_value == {'output': 'done\n'}
+
     async def test_run_code_syntax_error_becomes_model_retry(self) -> None:
         """A Python syntax error is surfaced as `ModelRetry` so the model can fix it."""
         wrapper = CodeMode[object]().get_wrapper_toolset(_build_function_toolset(add))