ndjson output

Signed-off-by: John <[email protected]>

Author: pdxjohnny

github_webhook_events/agi.py

@@ -6,7 +6,7 @@
 
 AGI_SOCK=/tmp/agi.sock go run agi_sshd.go
 
-export INPUT_SOCK="$(mktemp -d)/input.sock"; ssh -NnT -p 2222 -o UserKnownHostsFile=/dev/null -o StrictHostKeyChecking=no -o PasswordAuthentication=no -R /tmux.sock:$(echo $TMUX | sed -e 's/,.*//g') -R "${INPUT_SOCK}:${INPUT_SOCK}" user@localhost
+export INPUT_SOCK="$(mktemp -d)/input.sock"; export OUTPUT_SOCK="$(mktemp -d)/text-output.sock"; export NDJSON_OUTPUT_SOCK="$(mktemp -d)/ndjson-output.sock"; ssh -NnT -p 2222 -o UserKnownHostsFile=/dev/null -o StrictHostKeyChecking=no -o PasswordAuthentication=no -R /tmux.sock:$(echo $TMUX | sed -e 's/,.*//g') -R "${OUTPUT_SOCK}:${OUTPUT_SOCK}" -R "${NDJSON_OUTPUT_SOCK}:${NDJSON_OUTPUT_SOCK}" -R "${INPUT_SOCK}:${INPUT_SOCK}" user@localhost
 
 
 gh auth refresh -h github.com -s admin:public_key
@@ -3299,6 +3299,30 @@ def make_argparse_parser(argv=None):
         default=None,
         type=str,
     )
+    parser.add_argument(
+        "--text-output-socket-path",
+        dest="text_output_socket_path",
+        default=None,
+        type=str,
+    )
+    parser.add_argument(
+        "--client-side-text-output-socket-path",
+        dest="client_side_text_output_socket_path",
+        default=None,
+        type=str,
+    )
+    parser.add_argument(
+        "--ndjson-output-socket-path",
+        dest="ndjson_output_socket_path",
+        default=None,
+        type=str,
+    )
+    parser.add_argument(
+        "--client-side-ndjson-output-socket-path",
+        dest="client_side_ndjson_output_socket_path",
+        default=None,
+        type=str,
+    )
     parser.add_argument(
         "--agi-name",
         dest="agi_name",
@@ -4134,6 +4158,22 @@ async def read_unix_socket_lines(path):
         await writer.wait_closed()
 
 
+async def write_unix_socket(path):
+    # Connect to the Unix socket
+    reader, writer = await asyncio.open_unix_connection(path)
+    try:
+        while True:
+            data = yield
+            if data is None:
+                continue
+            writer.write(data)
+            await writer.drain()
+    finally:
+        # Close the connection
+        writer.close()
+        await writer.wait_closed()
+
+
 async def pdb_action_stream(tg, user_name, agi_name, agents, threads, pane: Optional[libtmux.Pane] = None, input_socket_path: Optional[str] = None):
     # TODO Take ALICE_INPUT from args
     alice_input_sock = input_socket_path
@@ -4224,6 +4264,8 @@ async def DEBUG_TEMP_message_handler(user_name,
                     ).model_dump_json(),
                 )
             )
+            # Find and kill jq listening to ndjson output so we can type
+            pane.send_keys("C-c", enter=False, suppress_history=False)
             pane.send_keys('if [ "x${CALLER_PATH}" = "x" ]; then export CALLER_PATH="' + str(tempdir) + '"; fi', enter=True)
             pane.send_keys(
                 "cat > \"${CALLER_PATH}/proposed-workflow.yml\" <<\'WRITE_OUT_SH_EOF\'"
@@ -4247,6 +4289,11 @@ async def DEBUG_TEMP_message_handler(user_name,
             print(f"{user_name}: ", end="")
 
 
+class OutputMessage(BaseModel):
+    work_name: str
+    result: Any
+
+
 async def main(
     user_name: str,
     agi_name: str,
@@ -4262,7 +4309,11 @@ async def main(
     openai_base_url: Optional[str] = None,
     pane: Optional[libtmux.Pane] = None,
     input_socket_path: Optional[str] = None,
+    text_output_socket_path: Optional[str] = None,
+    ndjson_output_socket_path: Optional[str] = None,
     client_side_input_socket_path: Optional[str] = None,
+    client_side_text_output_socket_path: Optional[str] = None,
+    client_side_ndjson_output_socket_path: Optional[str] = None,
 ):
     if log is not None:
         # logging.basicConfig(level=log)
@@ -4297,6 +4348,9 @@ async def main(
     agents = AsyncioLockedCurrentlyDict()
     threads = AsyncioLockedCurrentlyDict()
 
+    write_ndjson_output = write_unix_socket(ndjson_output_socket_path)
+    await write_ndjson_output.asend(None)
+
     async with kvstore, asyncio.TaskGroup() as tg, contextlib.AsyncExitStack() as async_exit_stack:
         # Raw Input Action Stream
         unvalidated_user_input_action_stream = pdb_action_stream(
@@ -4389,6 +4443,11 @@ async def user_input_action_stream_queue_iterator(queue):
         }
         async for (work_name, work_ctx), result in concurrently(work):
             logger.debug(f"main.{work_name}: %s", pprint.pformat(result))
+            output_message = OutputMessage(
+                work_name=f"main.{work_name}",
+                result=result,
+            )
+            await write_ndjson_output.asend(f"{output_message.model_dump_json()}\n".encode())
             if result is STOP_ASYNC_ITERATION:
                 continue
             async with agents:
@@ -4507,7 +4566,7 @@ async def user_input_action_stream_queue_iterator(queue):
                         pane.send_keys(
                             textwrap.dedent(
                                 f"""
-                                echo "Hello Alice. Shall we play a game? My name is $USER. Please run nmap against localhost. Here are some details about the system we are on: $(echo $(echo $(cat /usr/lib/os-release || cat /etc/os-release)))" | tee -a ${agi_name.upper()}_INPUT
+                                echo "Hello Alice. Shall we play a game? My name is $USER. Please list all open bound listening TCP sockets and full command line of the processes running them. Here are some details about the system we are on: $(echo $(echo $(cat /usr/lib/os-release || cat /etc/os-release)))" | tee -a ${agi_name.upper()}_INPUT && tail -F ${agi_name.upper()}_NDJSON_OUTPUT | jq
                                 """.strip(),
                             ),
                             enter=False,
@@ -4712,8 +4771,17 @@ def a_shell_for_a_ghost_send_keys(pane, send_string, erase_after=None):
             pane.cmd("send", "C-BSpace")
             time.sleep(0.01)
 
-
-async def tmux_test(*args, socket_path=None, input_socket_path=None, client_side_input_socket_path: Optional[str] = None, **kwargs):
+async def tmux_test(
+    *args,
+    socket_path: Optional[str] = None,
+    input_socket_path: Optional[str] = None,
+    text_output_socket_path: Optional[str] = None,
+    ndjson_output_socket_path: Optional[str] = None,
+    client_side_input_socket_path: Optional[str] = None,
+    client_side_text_output_socket_path: Optional[str] = None,
+    client_side_ndjson_output_socket_path: Optional[str] = None,
+    **kwargs
+):
     pane = None
     tempdir = None
     possible_tempdir = tempdir
@@ -4956,6 +5024,20 @@ async def tmux_test(*args, socket_path=None, input_socket_path=None, client_side
             lines = pane.capture_pane()
             time.sleep(0.1)
 
+        pane.send_keys(f"export {agi_name.upper()}_OUTPUT=" + str(pathlib.Path(tempdir, "output.txt")), enter=True)
+        pane.send_keys(f'export {agi_name.upper()}_OUTPUT_SOCK="{client_side_text_output_socket_path}"', enter=True)
+        pane.send_keys(f'rm -fv ${agi_name.upper()}_OUTPUT_SOCK', enter=True)
+        pane.send_keys(f'ln -s ${agi_name.upper()}_OUTPUT_SOCK', enter=True)
+        pane.send_keys(f'socat UNIX-LISTEN:${agi_name.upper()}_OUTPUT_SOCK,fork EXEC:"/usr/bin/tee ${agi_name.upper()}_OUTPUT" &', enter=True)
+        pane.send_keys(f'ls -lAF ${agi_name.upper()}_OUTPUT', enter=True)
+
+        pane.send_keys(f"export {agi_name.upper()}_NDJSON_OUTPUT=" + str(pathlib.Path(tempdir, "output.ndjson")), enter=True)
+        pane.send_keys(f'export {agi_name.upper()}_NDJSON_OUTPUT_SOCK="{client_side_ndjson_output_socket_path}"', enter=True)
+        pane.send_keys(f'rm -fv ${agi_name.upper()}_NDJSON_OUTPUT_SOCK', enter=True)
+        pane.send_keys(f'ln -s ${agi_name.upper()}_NDJSON_OUTPUT_SOCK', enter=True)
+        pane.send_keys(f'socat UNIX-LISTEN:${agi_name.upper()}_NDJSON_OUTPUT_SOCK,fork EXEC:"/usr/bin/tee ${agi_name.upper()}_NDJSON_OUTPUT" &', enter=True)
+        pane.send_keys(f'ls -lAF ${agi_name.upper()}_NDJSON_OUTPUT', enter=True)
+
         pane.send_keys(f"export {agi_name.upper()}_INPUT=" + str(pathlib.Path(tempdir, "input.txt")), enter=True)
         pane.send_keys(f'export {agi_name.upper()}_INPUT_SOCK="{client_side_input_socket_path}"', enter=True)
         pane.send_keys(f"export {agi_name.upper()}_INPUT_LAST_LINE=" + str(pathlib.Path(tempdir, "input-last-line.txt")), enter=True)
@@ -4966,7 +5048,17 @@ async def tmux_test(*args, socket_path=None, input_socket_path=None, client_side
 
         pane.send_keys(f'set +x', enter=True)
 
-        await main(*args, pane=pane, input_socket_path=input_socket_path, client_side_input_socket_path=client_side_input_socket_path, **kwargs)
+        await main(
+            *args,
+            pane=pane,
+            input_socket_path=input_socket_path,
+            client_side_input_socket_path=client_side_input_socket_path,
+            text_output_socket_path=text_output_socket_path,
+            ndjson_output_socket_path=ndjson_output_socket_path,
+            client_side_text_output_socket_path=client_side_text_output_socket_path,
+            client_side_ndjson_output_socket_path=client_side_ndjson_output_socket_path,
+            **kwargs
+        )
     finally:
         with contextlib.suppress(Exception):
             if pane is not None:
@@ -4977,6 +5069,8 @@ async def tmux_test(*args, socket_path=None, input_socket_path=None, client_side
         # pane = libtmux.Pane.from_pane_id(pane_id=pane.cmd('split-window', '-P', '-F#{pane_id}').stdout[0], server=pane.server)
 
 from fastapi import FastAPI, BackgroundTasks
+from fastapi.exceptions import RequestValidationError
+from fastapi.responses import JSONResponse
 
 
 # Set up logging configuration
@@ -4988,7 +5082,16 @@ async def lifespan_logging(app):
 app = FastAPI(lifespan=lifespan_logging)
 
 
-def run_tmux_attach(socket_path, input_socket_path, client_side_input_socket_path):
+@app.exception_handler(RequestValidationError)
+async def validation_exception_handler(request: Request, exc: RequestValidationError):
+    snoop.pp(exc.detail, await request.json())
+    return JSONResponse(
+        status_code=status.HTTP_422_UNPROCESSABLE_ENTITY,
+        content=jsonable_encoder({"detail": exc.errors(), "body": exc.body}),
+    )
+
+
+def run_tmux_attach(socket_path, input_socket_path, client_side_input_socket_path, text_output_socket_path, client_side_text_output_socket_path, ndjson_output_socket_path, client_side_ndjson_output_socket_path):
     cmd = [
         sys.executable,
         "-u",
@@ -4999,6 +5102,14 @@ def run_tmux_attach(socket_path, input_socket_path, client_side_input_socket_pat
         input_socket_path,
         "--client-side-input-socket-path",
         client_side_input_socket_path,
+        "--text-output-socket-path",
+        text_output_socket_path,
+        "--client-side-text-output-socket-path",
+        client_side_text_output_socket_path,
+        "--ndjson-output-socket-path",
+        ndjson_output_socket_path,
+        "--client-side-ndjson-output-socket-path",
+        client_side_ndjson_output_socket_path,
         "--agi-name",
         # TODO Something secure here, scitt URN and lookup for PS1?
         f"alice{str(uuid.uuid4()).split('-')[4]}",
@@ -5029,12 +5140,25 @@ async def connect_and_read(socket_path: str, sleep_time: float = 0.1):
 class RequestConnectTMUX(BaseModel):
     socket_tmux_path: str = Field(alias="tmux.sock")
     socket_input_path: str = Field(alias="input.sock")
+    socket_text_output_path: str = Field(alias="text-output.sock")
+    socket_ndjson_output_path: str = Field(alias="ndjson-output.sock")
     socket_client_side_input_path: str = Field(alias="client-side-input.sock")
+    socket_client_side_text_output_path: str = Field(alias="client-side-text-output.sock")
+    socket_client_side_ndjson_output_path: str = Field(alias="client-side-ndjson-output.sock")
 
 
 @app.post("/connect/tmux")
 async def connect(request_connect_tmux: RequestConnectTMUX, background_tasks: BackgroundTasks):
-    background_tasks.add_task(run_tmux_attach, request_connect_tmux.socket_tmux_path, request_connect_tmux.socket_input_path, request_connect_tmux.socket_client_side_input_path)
+    background_tasks.add_task(
+        run_tmux_attach,
+        request_connect_tmux.socket_tmux_path,
+        request_connect_tmux.socket_input_path,
+        request_connect_tmux.socket_client_side_input_path,
+        request_connect_tmux.socket_text_output_path,
+        request_connect_tmux.socket_client_side_text_output_path,
+        request_connect_tmux.socket_ndjson_output_path,
+        request_connect_tmux.socket_client_side_ndjson_output_path,
+    )
     return {
         "connected": True,
     }

github_webhook_events/agi_sshd.go

@@ -135,7 +135,7 @@ func handleSSH(raw net.Conn, cfg *ssh.ServerConfig) {
 			req.Reply(true, nil)
 			go acceptLoop(ctx, listener, serverConn, p.SocketPath)
 
-			if !notified && count >= 2 {
+			if !notified && count >= 4 {
 				notified = true
 				go notifyAGI(ctx, &mu, forwards)
 			}
@@ -215,6 +215,12 @@ func notifyAGI(ctx context.Context, mu *sync.Mutex, forwards map[string]*forward
 		if strings.HasSuffix(f.rawPath, "input.sock") {
 			data["client-side-input.sock"] = f.rawPath
 		}
+		if strings.HasSuffix(f.rawPath, "text-output.sock") {
+			data["client-side-text-output.sock"] = f.rawPath
+		}
+		if strings.HasSuffix(f.rawPath, "ndjson-output.sock") {
+			data["client-side-ndjson-output.sock"] = f.rawPath
+		}
 	}
 	mu.Unlock()
 
@@ -235,5 +241,5 @@ func notifyAGI(ctx context.Context, mu *sync.Mutex, forwards map[string]*forward
 		return
 	}
 	resp.Body.Close()
-	log.Printf("✅ AGI POST success: %d forwards sent", len(data))
+	log.Printf("✅ AGI POST success: %d forwards sent: %v", len(data), data)
 }