535 lines
16 KiB
Go
535 lines
16 KiB
Go
// ARTEMIS adapter for the pinned MIT Cua Driver. No model, scheduler, or recorder.
|
|
package main
|
|
|
|
import (
|
|
"bufio"
|
|
"bytes"
|
|
"crypto/sha256"
|
|
_ "embed"
|
|
"encoding/base64"
|
|
"encoding/json"
|
|
"errors"
|
|
"fmt"
|
|
"image"
|
|
_ "image/jpeg"
|
|
_ "image/png"
|
|
"io"
|
|
"os"
|
|
"os/exec"
|
|
"path/filepath"
|
|
"strings"
|
|
"sync"
|
|
"time"
|
|
)
|
|
|
|
//go:embed tools.json
|
|
var toolsJSON []byte
|
|
|
|
const adapterVersion = "0.36.0+artemis.4"
|
|
|
|
type obj = map[string]any
|
|
type request struct {
|
|
JSONRPC string `json:"jsonrpc"`
|
|
ID json.RawMessage `json:"id"`
|
|
Method string `json:"method"`
|
|
Params obj `json:"params"`
|
|
}
|
|
type peer struct {
|
|
cmd *exec.Cmd
|
|
input io.WriteCloser
|
|
lines chan []byte
|
|
next int
|
|
mu sync.Mutex
|
|
dead bool
|
|
job uintptr
|
|
}
|
|
|
|
func object(v any) obj {
|
|
m, _ := v.(map[string]any)
|
|
if m == nil {
|
|
return obj{}
|
|
}
|
|
return m
|
|
}
|
|
func str(v any) string { s, _ := v.(string); return s }
|
|
func num(v any) int64 {
|
|
switch n := v.(type) {
|
|
case float64:
|
|
return int64(n)
|
|
case int64:
|
|
return n
|
|
case int:
|
|
return int64(n)
|
|
}
|
|
return 0
|
|
}
|
|
func items(v any) []any { a, _ := v.([]any); return a }
|
|
func textBlock(s string) obj { return obj{"type": "text", "text": s} }
|
|
func receipt(execution, side, code string) obj {
|
|
return obj{"version": 1, "execution": execution, "side_effect": side, "retryable": false, "error_code": code}
|
|
}
|
|
func failure(msg, execution, side, code string) obj {
|
|
return obj{"isError": true, "content": []any{textBlock(msg)}, "structuredContent": obj{"artemis": receipt(execution, side, code)}}
|
|
}
|
|
func writeTool(name string) bool {
|
|
return name != "list_windows" && name != "get_window_state" && name != "verify_state"
|
|
}
|
|
|
|
func startPeer(exe string) (*peer, error) {
|
|
c := exec.Command(exe, "mcp", "--direct")
|
|
c.Env = append(os.Environ(), "DO_NOT_TRACK=1", "CUA_TELEMETRY=0", "CUA_DRIVER_RS_TELEMETRY_ENABLED=false")
|
|
configureProcess(c)
|
|
in, err := c.StdinPipe()
|
|
if err != nil {
|
|
return nil, err
|
|
}
|
|
out, err := c.StdoutPipe()
|
|
if err != nil {
|
|
return nil, err
|
|
}
|
|
c.Stderr = os.Stderr
|
|
if err = c.Start(); err != nil {
|
|
return nil, err
|
|
}
|
|
job, err := ownProcess(c.Process.Pid)
|
|
if err != nil {
|
|
_ = c.Process.Kill()
|
|
_ = c.Wait()
|
|
return nil, err
|
|
}
|
|
p := &peer{cmd: c, input: in, lines: make(chan []byte, 32), job: job}
|
|
go func() {
|
|
defer close(p.lines)
|
|
s := bufio.NewScanner(out)
|
|
s.Buffer(make([]byte, 65536), 32<<20)
|
|
for s.Scan() {
|
|
p.lines <- append([]byte(nil), s.Bytes()...)
|
|
}
|
|
}()
|
|
_, err = p.call("initialize", obj{"protocolVersion": "2024-11-05", "clientInfo": obj{"name": "artemis-cua-adapter", "version": adapterVersion}, "capabilities": obj{}})
|
|
if err != nil {
|
|
p.stop()
|
|
return nil, err
|
|
}
|
|
err = json.NewEncoder(in).Encode(obj{"jsonrpc": "2.0", "method": "notifications/initialized"})
|
|
if err != nil {
|
|
p.stop()
|
|
return nil, err
|
|
}
|
|
return p, nil
|
|
}
|
|
func (p *peer) stop() {
|
|
p.mu.Lock()
|
|
defer p.mu.Unlock()
|
|
if p.dead {
|
|
return
|
|
}
|
|
p.dead = true
|
|
_ = p.input.Close()
|
|
closeJob(p.job)
|
|
_ = p.cmd.Process.Kill()
|
|
_ = p.cmd.Wait()
|
|
}
|
|
func (p *peer) call(method string, params obj) (obj, error) {
|
|
p.mu.Lock()
|
|
dead := p.dead
|
|
p.mu.Unlock()
|
|
if dead {
|
|
return nil, errors.New("driver closed")
|
|
}
|
|
p.next++
|
|
id := p.next
|
|
if err := json.NewEncoder(p.input).Encode(obj{"jsonrpc": "2.0", "id": id, "method": method, "params": params}); err != nil {
|
|
return nil, err
|
|
}
|
|
timer := time.NewTimer(45 * time.Second)
|
|
defer timer.Stop()
|
|
for {
|
|
select {
|
|
case raw, ok := <-p.lines:
|
|
if !ok {
|
|
return nil, errors.New("driver disconnected")
|
|
}
|
|
var r obj
|
|
if json.Unmarshal(raw, &r) != nil {
|
|
continue
|
|
}
|
|
if num(r["id"]) != int64(id) {
|
|
continue
|
|
}
|
|
if r["error"] != nil {
|
|
return nil, errors.New("driver protocol error")
|
|
}
|
|
return object(r["result"]), nil
|
|
case <-timer.C:
|
|
p.stop()
|
|
return nil, errors.New("driver response timed out; action not replayed")
|
|
}
|
|
}
|
|
}
|
|
func (p *peer) tool(name string, args obj) (obj, error) {
|
|
return p.call("tools/call", obj{"name": name, "arguments": args})
|
|
}
|
|
|
|
type adapter struct {
|
|
peer *peer
|
|
tools []any
|
|
schemas map[string]obj
|
|
captures map[string]pointerFrame
|
|
}
|
|
|
|
func newAdapter(p *peer) *adapter {
|
|
a := &adapter{peer: p, schemas: map[string]obj{}, captures: map[string]pointerFrame{}}
|
|
_ = json.Unmarshal(toolsJSON, &a.tools)
|
|
for _, v := range a.tools {
|
|
t := object(v)
|
|
a.schemas[str(t["name"])] = object(t["inputSchema"])
|
|
}
|
|
return a
|
|
}
|
|
func (a *adapter) validate(name string, args obj) error {
|
|
schema, ok := a.schemas[name]
|
|
if !ok {
|
|
return errors.New("tool is not exposed")
|
|
}
|
|
if name == "press_key" || name == "type_text" {
|
|
for _, key := range []string{"element_token", "x", "y"} {
|
|
if _, supplied := args[key]; supplied {
|
|
return errors.New(name + " uses the window's current focus; omit element_token/x/y. Focus with click/click_at first, or use set_value for a UIA field")
|
|
}
|
|
}
|
|
}
|
|
props := object(schema["properties"])
|
|
for key := range args {
|
|
if _, ok := props[key]; !ok {
|
|
return fmt.Errorf("unsupported argument: %s", key)
|
|
}
|
|
}
|
|
for _, key := range items(schema["required"]) {
|
|
if _, ok := args[str(key)]; !ok {
|
|
return fmt.Errorf("missing argument: %s", str(key))
|
|
}
|
|
}
|
|
if name == "list_windows" {
|
|
_, err := windowQuery(args)
|
|
return err
|
|
}
|
|
if num(args["pid"]) <= 0 || num(args["window_id"]) <= 0 {
|
|
return errors.New("a concrete pid and window_id are required")
|
|
}
|
|
if name == "set_value" && strings.TrimSpace(str(args["element_token"])) == "" {
|
|
return errors.New("set_value requires the latest element_token; use type_text for the current focus")
|
|
}
|
|
if name == "save_screenshot" && !validFilename(str(args["filename"])) {
|
|
return errors.New("filename must be a new plain PNG filename")
|
|
}
|
|
if name == "get_window_state" && (args["include_screenshot"] == false || args["include_accessibility_tree"] == false) {
|
|
return errors.New("candidate observations require both screenshot and accessibility tree")
|
|
}
|
|
return nil
|
|
}
|
|
func validFilename(s string) bool {
|
|
return s != "" && len(s) <= 160 && strings.EqualFold(filepath.Ext(s), ".png") && !strings.ContainsAny(s, "/\\:<>\"|?*\x00") && s != "." && s != ".." && !strings.HasSuffix(s, ".") && !strings.HasSuffix(s, " ")
|
|
}
|
|
func actionReceipt(r obj, write bool) obj {
|
|
if !write {
|
|
if r["isError"] == true {
|
|
return receipt("not_started", "none", "observation_failed")
|
|
}
|
|
return receipt("completed", "none", "")
|
|
}
|
|
s := object(r["structuredContent"])
|
|
// v0.36.0 class-level preflight refusal has no effect field. It occurs
|
|
// before input injection. Do not classify actuator failures with a cause
|
|
// as safe-to-repeat; retain their uncertain execution below.
|
|
if s["code"] == "background_unavailable" && s["cause"] == nil && s["event_kind"] != nil && r["isError"] == true {
|
|
return receipt("not_started", "none", "background_unavailable")
|
|
}
|
|
switch str(s["effect"]) {
|
|
case "confirmed":
|
|
return receipt("completed", "confirmed", "")
|
|
case "partial":
|
|
return receipt("partial", "confirmed", "partial_delivery")
|
|
case "refused":
|
|
return receipt("not_started", "none", "driver_refused")
|
|
case "unverifiable", "suspected_noop":
|
|
return receipt("completed", "possible", str(s["effect"]))
|
|
default:
|
|
return receipt("unknown", "possible", "unclassified_result")
|
|
}
|
|
}
|
|
func imageSize(r obj) (int, int) {
|
|
for _, v := range items(r["content"]) {
|
|
c := object(v)
|
|
if c["type"] == "image" {
|
|
raw, err := base64.StdEncoding.DecodeString(str(c["data"]))
|
|
if err == nil {
|
|
cfg, _, err := image.DecodeConfig(bytes.NewReader(raw))
|
|
if err == nil {
|
|
return cfg.Width, cfg.Height
|
|
}
|
|
}
|
|
}
|
|
}
|
|
return 0, 0
|
|
}
|
|
func source(r obj, args obj, started float64) obj {
|
|
s := object(r["structuredContent"])
|
|
w, h := imageSize(r)
|
|
if r["isError"] == true || w <= 0 || h <= 0 || num(s["pid"]) != num(args["pid"]) || num(s["window_id"]) != num(args["window_id"]) {
|
|
return nil
|
|
}
|
|
bounds := object(s["window_bounds"])
|
|
rect := obj{"left": num(bounds["x"]), "top": num(bounds["y"]), "width": num(bounds["width"]), "height": num(bounds["height"])}
|
|
screen, ok := verifyWindow(num(args["pid"]), num(args["window_id"]), rect)
|
|
if !ok {
|
|
return nil
|
|
}
|
|
return obj{"kind": "computer", "target_id": fmt.Sprint(num(args["window_id"])), "label": s["window_title"], "captured_at": started, "width": w, "height": h, "window": rect, "screen": screen}
|
|
}
|
|
func (a *adapter) observe(args obj) (obj, obj, error) {
|
|
delete(a.captures, targetKey(args))
|
|
started := float64(time.Now().UnixMilli()) / 1000
|
|
before := foregroundWindow()
|
|
r, err := a.peer.tool("get_window_state", args)
|
|
if err != nil {
|
|
return nil, nil, err
|
|
}
|
|
annotateCaptureFocus(r, num(args["window_id"]), before, foregroundWindow())
|
|
src := source(r, args, started)
|
|
if frame, ok := capturedFrame(r, src); ok {
|
|
a.captures[targetKey(args)] = frame
|
|
prependFrameHint(r, frame)
|
|
}
|
|
return r, src, nil
|
|
}
|
|
func observationArgs(args obj) obj {
|
|
return obj{"pid": args["pid"], "window_id": args["window_id"], "max_image_dimension": 1280, "max_elements": 250, "tree_format": "markdown"}
|
|
}
|
|
func (a *adapter) handle(name string, args obj) obj {
|
|
if err := a.validate(name, args); err != nil {
|
|
return failure(err.Error(), "not_started", "none", "invalid_input")
|
|
}
|
|
if name != "list_windows" {
|
|
if !windowOwned(num(args["pid"]), num(args["window_id"])) {
|
|
return failure("window is gone or belongs to a different process", "not_started", "none", "invalid_target")
|
|
}
|
|
}
|
|
if name == "save_screenshot" {
|
|
return a.save(args)
|
|
}
|
|
var r, src obj
|
|
var err error
|
|
if name == "get_window_state" {
|
|
for k, v := range observationArgs(args) {
|
|
if _, ok := args[k]; !ok {
|
|
args[k] = v
|
|
}
|
|
}
|
|
r, src, err = a.observe(args)
|
|
} else if name == "list_windows" {
|
|
query, _ := windowQuery(args) // Already validated above.
|
|
r, err = a.peer.tool(name, query)
|
|
if err == nil {
|
|
r = windowListing(r, query)
|
|
}
|
|
} else {
|
|
if isPointerTool(name) {
|
|
frame, ok := a.captures[targetKey(args)]
|
|
if !ok {
|
|
return failure("No current screenshot; get_window_state before a pixel action", "not_started", "none", "invalid_coordinates")
|
|
}
|
|
if err := frame.validate(name, args); err != nil {
|
|
return failure(err.Error(), "not_started", "none", "invalid_coordinates")
|
|
}
|
|
if _, ok := verifyWindow(num(args["pid"]), num(args["window_id"]), frame.window); !ok {
|
|
delete(a.captures, targetKey(args))
|
|
return failure("Window moved or resized; get_window_state again", "not_started", "none", "stale_capture")
|
|
}
|
|
if name == "scroll" {
|
|
// v0.36.0 scroll omits screenshot scaling. Measure actual native
|
|
// capture pixels (not DWM/DPI estimates) without another model call.
|
|
// Do not expose this internal frame or reuse its tokens.
|
|
delete(a.captures, targetKey(args))
|
|
prepared, e := nativeScrollArgs(frame, args, func() (pointerFrame, error) {
|
|
nativeArgs := observationArgs(args)
|
|
nativeArgs["max_image_dimension"] = 0
|
|
nativeArgs["include_accessibility_tree"] = false
|
|
native, err := a.peer.tool("get_window_state", nativeArgs)
|
|
if err != nil {
|
|
return pointerFrame{}, err
|
|
}
|
|
src := source(native, args, float64(time.Now().UnixMilli())/1000)
|
|
f, ok := capturedFrame(native, src)
|
|
if !ok {
|
|
return pointerFrame{}, errors.New("native capture unavailable; get_window_state again")
|
|
}
|
|
return f, nil
|
|
})
|
|
if e != nil {
|
|
return failure(e.Error(), "not_started", "none", "invalid_coordinates")
|
|
}
|
|
args = prepared
|
|
} else if name == "drag" {
|
|
// Upstream drag already scales using its current screenshot.
|
|
args = copyArgs(args)
|
|
delete(args, "capture_id")
|
|
}
|
|
}
|
|
upstream := name
|
|
if name == "click_at" {
|
|
upstream = "click"
|
|
}
|
|
if name == "verify_state" {
|
|
delete(a.captures, targetKey(args))
|
|
}
|
|
if writeTool(name) {
|
|
delete(a.captures, targetKey(args))
|
|
}
|
|
r, err = callWithForeground(a.peer.tool, upstream, args, foregroundWindow)
|
|
}
|
|
if err != nil {
|
|
a.peer.stop()
|
|
if writeTool(name) {
|
|
return failure(err.Error(), "unknown", "possible", "transport_lost")
|
|
}
|
|
return failure(err.Error(), "not_started", "none", "transport_lost")
|
|
}
|
|
structured := object(r["structuredContent"])
|
|
rec := object(structured["artemis"])
|
|
if len(rec) == 0 {
|
|
rec = actionReceipt(r, writeTool(name))
|
|
}
|
|
if writeTool(name) && rec["execution"] != "not_started" {
|
|
// Observe once after delivery; NEVER repeat an uncertain write.
|
|
observation, origin, e := a.observe(observationArgs(args))
|
|
if e == nil && origin != nil {
|
|
src = origin
|
|
structured["observation"] = observation["structuredContent"]
|
|
r["content"] = append(items(r["content"]), items(observation["content"])...)
|
|
} else {
|
|
structured["observation_error"] = "Post-action evidence unavailable; inspect before another action."
|
|
}
|
|
}
|
|
if name == "verify_state" { // Read-only verification may include a screenshot; it is not a task verdict.
|
|
src = source(r, args, float64(time.Now().UnixMilli())/1000)
|
|
if frame, ok := capturedFrame(r, src); ok {
|
|
a.captures[targetKey(args)] = frame
|
|
prependFrameHint(r, frame)
|
|
}
|
|
}
|
|
if src != nil {
|
|
rec["source"] = src
|
|
} else if name == "get_window_state" && r["isError"] != true {
|
|
return failure("capture target or geometry could not be validated", "not_started", "none", "invalid_capture")
|
|
}
|
|
if structured["effect"] == "refused" {
|
|
r["isError"] = true
|
|
}
|
|
structured["artemis"] = rec
|
|
r["structuredContent"] = structured
|
|
return r
|
|
}
|
|
func (a *adapter) save(args obj) obj {
|
|
root := os.Getenv("ARTEMIS_ARTIFACTS_DIR")
|
|
if root == "" {
|
|
return failure("task artifact directory unavailable", "not_started", "none", "missing_artifacts")
|
|
}
|
|
r, src, err := a.observe(observationArgs(args))
|
|
if err != nil || src == nil {
|
|
return failure("fresh window capture failed", "not_started", "none", "capture_failed")
|
|
}
|
|
var png []byte
|
|
for _, v := range items(r["content"]) {
|
|
b := object(v)
|
|
if b["type"] == "image" && b["mimeType"] == "image/png" {
|
|
png, _ = base64.StdEncoding.DecodeString(str(b["data"]))
|
|
break
|
|
}
|
|
}
|
|
if len(png) == 0 {
|
|
return failure("expected PNG capture", "not_started", "none", "capture_format")
|
|
}
|
|
if err = os.MkdirAll(root, 0700); err != nil {
|
|
return failure("artifact directory unavailable", "not_started", "none", "file_error")
|
|
}
|
|
path := filepath.Join(root, str(args["filename"]))
|
|
f, err := os.OpenFile(path, os.O_WRONLY|os.O_CREATE|os.O_EXCL, 0600)
|
|
if err != nil {
|
|
return failure("file exists or cannot be created", "not_started", "none", "file_error")
|
|
}
|
|
n, writeErr := f.Write(png)
|
|
closeErr := f.Close()
|
|
actual, readErr := os.ReadFile(path)
|
|
if writeErr != nil || closeErr != nil || readErr != nil || n != len(png) || !bytes.Equal(actual, png) {
|
|
return failure("artifact write could not be verified", "partial", "possible", "file_error")
|
|
}
|
|
rec := receipt("completed", "confirmed", "")
|
|
rec["source"] = src
|
|
rec["attachments"] = []any{obj{"name": str(args["filename"]), "sha256": fmt.Sprintf("%x", sha256.Sum256(actual)), "media_type": "image/png"}}
|
|
object(r["structuredContent"])["artemis"] = rec
|
|
r["content"] = append(items(r["content"]), textBlock("Saved verified task attachment: "+str(args["filename"])))
|
|
return r
|
|
}
|
|
func serve(a *adapter, in io.Reader, out io.Writer) {
|
|
input := make(chan request)
|
|
go func() {
|
|
defer close(input)
|
|
defer a.peer.stop()
|
|
s := bufio.NewScanner(in)
|
|
s.Buffer(make([]byte, 65536), 2<<20)
|
|
for s.Scan() {
|
|
var q request
|
|
if json.Unmarshal(s.Bytes(), &q) == nil {
|
|
if q.Method == "notifications/cancelled" {
|
|
a.peer.stop()
|
|
continue
|
|
}
|
|
input <- q
|
|
}
|
|
}
|
|
}()
|
|
enc := json.NewEncoder(out)
|
|
for q := range input {
|
|
if len(q.ID) == 0 {
|
|
continue
|
|
}
|
|
var result any
|
|
var rpcErr any
|
|
switch q.Method {
|
|
case "initialize":
|
|
result = obj{"protocolVersion": "2024-11-05", "capabilities": obj{"tools": obj{}}, "serverInfo": obj{"name": "artemis-cua-driver", "version": adapterVersion}}
|
|
case "ping":
|
|
result = obj{}
|
|
case "tools/list":
|
|
result = obj{"tools": a.tools}
|
|
case "tools/call":
|
|
result = a.handle(str(q.Params["name"]), object(q.Params["arguments"]))
|
|
default:
|
|
rpcErr = obj{"code": -32601, "message": "method not found"}
|
|
}
|
|
response := obj{"jsonrpc": "2.0", "id": q.ID}
|
|
if rpcErr != nil {
|
|
response["error"] = rpcErr
|
|
} else {
|
|
response["result"] = result
|
|
}
|
|
if enc.Encode(response) != nil {
|
|
return
|
|
}
|
|
}
|
|
}
|
|
func main() {
|
|
exe, err := os.Executable()
|
|
if err != nil {
|
|
os.Exit(1)
|
|
}
|
|
p, err := startPeer(filepath.Join(filepath.Dir(exe), "cua-driver.exe"))
|
|
if err != nil {
|
|
fmt.Fprintln(os.Stderr, "Cua startup failed:", err)
|
|
os.Exit(1)
|
|
}
|
|
defer p.stop()
|
|
serve(newAdapter(p), os.Stdin, os.Stdout)
|
|
}
|