personal memory agent
1# SPDX-License-Identifier: AGPL-3.0-only
2# Copyright (c) 2026 sol pbc
3
4import asyncio
5import importlib
6import json
7import sys
8from unittest.mock import AsyncMock
9
10from tests.conftest import setup_google_genai_stub
11from tests.test_google import make_mock_process
12from think.models import GEMINI_FLASH
13
14
15async def run_main(mod, argv, stdin_data=None):
16 sys.argv = argv
17 if stdin_data:
18 import io
19
20 sys.stdin = io.StringIO(stdin_data)
21 await mod.main_async()
22
23
24def test_google_thinking_events(monkeypatch, tmp_path, capsys):
25 setup_google_genai_stub(monkeypatch, with_thinking=True)
26
27 sys.modules.pop("think.providers.google", None)
28 importlib.reload(importlib.import_module("think.providers.google"))
29 mod = importlib.reload(importlib.import_module("think.agents"))
30
31 journal = tmp_path / "journal"
32 journal.mkdir()
33
34 monkeypatch.setenv("_SOLSTONE_JOURNAL_OVERRIDE", str(journal))
35 monkeypatch.setenv("GOOGLE_API_KEY", "x")
36 monkeypatch.setattr(
37 "think.providers.cli.shutil.which",
38 lambda name: "/usr/bin/gemini" if name == "gemini" else None,
39 )
40
41 # Simulate CLI output: thinking text before a tool call triggers thinking event
42 stdout_lines = [
43 json.dumps(
44 {
45 "type": "init",
46 "timestamp": 100,
47 "session_id": "sess-think",
48 "model": "gemini-2.5-flash",
49 }
50 ),
51 json.dumps(
52 {
53 "type": "message",
54 "role": "assistant",
55 "delta": True,
56 "content": "I need to analyze this step by step.",
57 }
58 ),
59 json.dumps(
60 {
61 "type": "tool_use",
62 "timestamp": 150,
63 "tool_name": "search_insights",
64 "tool_id": "t1",
65 "parameters": {"query": "hello"},
66 }
67 ),
68 json.dumps(
69 {
70 "type": "tool_result",
71 "timestamp": 200,
72 "tool_id": "t1",
73 "status": "success",
74 "output": "result data",
75 }
76 ),
77 json.dumps(
78 {
79 "type": "message",
80 "role": "assistant",
81 "delta": True,
82 "content": "ok",
83 }
84 ),
85 json.dumps(
86 {
87 "type": "result",
88 "timestamp": 300,
89 "status": "success",
90 "stats": {
91 "total_tokens": 20,
92 "input_tokens": 10,
93 "output_tokens": 10,
94 },
95 }
96 ),
97 ]
98 process = make_mock_process(stdout_lines)
99 monkeypatch.setattr(
100 "think.providers.cli.asyncio.create_subprocess_exec",
101 AsyncMock(return_value=process),
102 )
103
104 ndjson_input = json.dumps(
105 {
106 "prompt": "hello",
107 "provider": "google",
108 "model": GEMINI_FLASH,
109 "tools": ["search_insights"],
110 }
111 )
112 asyncio.run(run_main(mod, ["sol agents"], stdin_data=ndjson_input))
113
114 out_lines = capsys.readouterr().out.strip().splitlines()
115 events = [json.loads(line) for line in out_lines]
116
117 # Check that we have start, thinking, and finish events
118 assert events[0]["event"] == "start"
119 assert isinstance(events[0]["ts"], int)
120 assert "hello" in events[0]["prompt"]
121
122 # Look for thinking event (flushed when tool_use arrives)
123 thinking_events = [e for e in events if e["event"] == "thinking"]
124 assert len(thinking_events) == 1
125 assert thinking_events[0]["summary"] == "I need to analyze this step by step."
126 assert thinking_events[0]["model"] == GEMINI_FLASH
127 assert isinstance(thinking_events[0]["ts"], int)
128
129 assert events[-1]["event"] == "finish"
130 assert events[-1]["result"] == "ok"